brute 1.0.1 → 2.0.0

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

@@ -1,68 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    module OTel
-      # Records tool calls the LLM requested as span events.
-      #
-      # Runs POST-call: after the LLM responds, inspects ctx.functions
-      # for any tool calls the model wants to make, and adds a span event
-      # for each one with the tool name, call ID, and arguments.
-      #
-      class ToolCalls < Base
-        def call(env)
-          response = @app.call(env)
-
-          span = env[:span]
-          if span
-            functions = env[:pending_functions]
-            if functions && !functions.empty?
-              span.set_attribute("brute.tool_calls.count", functions.size)
-
-              functions.each do |fn|
-                attrs = {
-                  "tool.name" => fn.name.to_s,
-                  "tool.id" => fn.id.to_s,
-                }
-                args = fn.arguments
-                attrs["tool.arguments"] = args.to_json if args
-                span.add_event("tool_call", attributes: attrs)
-              end
-            end
-          end
-
-          response
-        end
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../../spec/support/mock_provider"
-  require_relative "../../../../spec/support/mock_response"
-
-  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 "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
-
-  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,65 +0,0 @@
-# frozen_string_literal: true
-
-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 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.
-      #
-      class ToolResults < Base
-        def call(env)
-          span = env[:span]
-
-          if span && (results = env[:tool_results])
-            span.set_attribute("brute.tool_results.count", results.size)
-
-            results.each do |name, value|
-              error = value.is_a?(Hash) && value[:error]
-              attrs = { "tool.name" => name.to_s }
-              if error
-                attrs["tool.status"] = "error"
-                attrs["tool.error"] = value[:error].to_s
-              else
-                attrs["tool.status"] = "ok"
-              end
-              span.add_event("tool_result", attributes: attrs)
-            end
-          end
-
-          @app.call(env)
-        end
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../../spec/support/mock_provider"
-  require_relative "../../../../spec/support/mock_response"
-
-  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 "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
-
-  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/otel.rb

@@ -1,34 +0,0 @@
-# frozen_string_literal: true
-
-if __FILE__ == $0
-  require "bundler/setup"
-  require "brute"
-end
-
-module Brute
-  module Middleware
-    # OpenTelemetry instrumentation for the LLM pipeline.
-    #
-    # Each middleware is independent and communicates through env[:span].
-    # OTel::Span must be outermost — it creates the span. The rest
-    # decorate it with events and attributes from their position in the
-    # pipeline.
-    #
-    # All middlewares are no-ops when opentelemetry-sdk is not loaded.
-    #
-    # Usage in pipeline:
-    #
-    #   use Brute::Middleware::OTel::Span
-    #   use Brute::Middleware::OTel::ToolResults
-    #   use Brute::Middleware::OTel::ToolCalls
-    #   use Brute::Middleware::OTel::TokenUsage
-    #
-    module OTel
-    end
-  end
-end
-
-require_relative "otel/span"
-require_relative "otel/tool_results"
-require_relative "otel/tool_calls"
-require_relative "otel/token_usage"

data/lib/brute/middleware/reasoning_normalizer.rb

@@ -1,192 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    # Handles reasoning/thinking content across model switches.
-    #
-    # PRE-call:
-    #   - If reasoning is enabled, injects provider-specific params into
-    #     the env (e.g., Anthropic thinking config, OpenAI reasoning_effort).
-    #   - Tracks which model produced each message. When the model changes,
-    #     strips reasoning_content from messages produced by the old model
-    #     (signatures are model-specific and cryptographically tied).
-    #
-    # POST-call:
-    #   - Records the current model on the response for future normalization.
-    #
-    # llm.rb exposes:
-    #   - response.reasoning_content  — the thinking text
-    #   - response.reasoning_tokens   — token count
-    #   - Provider params pass-through — we can send thinking:, reasoning_effort:, etc.
-    #
-    class ReasoningNormalizer < Base
-      # Effort levels that map to provider-specific params.
-      # Mirrors forgecode's Effort enum.
-      EFFORT_LEVELS = {
-        none: "none",
-        minimal: "low",
-        low: "low",
-        medium: "medium",
-        high: "high",
-        xhigh: "high",
-        max: "high",
-      }.freeze
-
-      def initialize(app, model_id: nil, effort: :medium, enabled: true, budget_tokens: nil)
-        super(app)
-        @model_id = model_id
-        @effort = effort
-        @enabled = enabled
-        @budget_tokens = budget_tokens
-        @message_models = [] # tracks which model produced each assistant message
-      end
-
-      def call(env)
-        if @enabled
-          inject_reasoning_params!(env)
-        end
-
-        response = @app.call(env)
-
-        # POST: record which model produced this response
-        if response
-          @message_models << @model_id
-        end
-
-        response
-      end
-
-      # Update the active model (e.g., when user switches models mid-session).
-      def model_id=(new_model)
-        @model_id = new_model
-      end
-
-      private
-
-      def inject_reasoning_params!(env)
-        env[:params] ||= {}
-        provider = env[:provider]
-
-        case provider_type(provider)
-        when :anthropic
-          if @budget_tokens
-            # Older extended thinking API (claude-3.7-sonnet style)
-            env[:params][:thinking] = {type: "enabled", budget_tokens: @budget_tokens}
-          else
-            # Newer effort-based API (claude-4 style) — pass through
-            # Anthropic handles this via the model itself
-          end
-        when :openai
-          env[:params][:reasoning_effort] = EFFORT_LEVELS[@effort] || "medium"
-        end
-      end
-
-      def provider_type(provider)
-        class_name = provider.class.name.to_s.downcase
-        if class_name.include?("anthropic")
-          :anthropic
-        elsif class_name.include?("openai")
-          :openai
-        elsif class_name.include?("google") || class_name.include?("gemini")
-          :google
-        else
-          :unknown
-        end
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../spec/support/mock_provider"
-  require_relative "../../../spec/support/mock_response"
-
-  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
-
-  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
-    klass.define_method(:class) do
-      c = super()
-      name_str = "LLM::#{type_name}"
-      c.define_singleton_method(:name) { name_str }
-      c
-    end
-    klass.new
-  end
-
-  inner_app = ->(_env) { MockResponse.new(content: "reasoned response") }
-
-  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
-
-  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
-
-  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
-
-  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 "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
-
-  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
-
-  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

data/lib/brute/middleware/retry.rb

@@ -1,157 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    # Retries the inner call on transient LLM errors with exponential backoff.
-    #
-    # Catches LLM::RateLimitError and LLM::ServerError, sleeps with
-    # exponential delay, and re-calls the inner app. Non-retryable errors
-    # propagate immediately.
-    #
-    # Unlike forgecode's separate retry.rs, this middleware wraps the LLM call
-    # directly — it sees the error and retries without the agent loop knowing.
-    #
-    class Retry < Base
-      DEFAULT_MAX_ATTEMPTS = 3
-      DEFAULT_BASE_DELAY = 2 # seconds
-
-      def initialize(app, max_attempts: DEFAULT_MAX_ATTEMPTS, base_delay: DEFAULT_BASE_DELAY)
-        super(app)
-        @max_attempts = max_attempts
-        @base_delay = base_delay
-      end
-
-      def call(env)
-        attempts = 0
-        begin
-          @app.call(env)
-        rescue LLM::RateLimitError, LLM::ServerError => e
-          attempts += 1
-          if attempts >= @max_attempts
-            env[:metadata][:last_error] = e.message
-            raise
-          end
-
-          delay = @base_delay ** attempts
-          env[:metadata][:retry_attempt] = attempts
-          env[:metadata][:retry_delay] = delay
-
-          sleep(delay)
-          retry
-        end
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../spec/support/mock_provider"
-  require_relative "../../../spec/support/mock_response"
-
-  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
-
-  def mock_inner_app(response:)
-    calls = []
-    app = ->(env) { calls << env; response }
-    [app, calls]
-  end
-
-  def flaky_inner_app(error_class, fail_count:, response:)
-    attempt = 0
-    ->(env) { attempt += 1; raise error_class, "transient" if attempt <= fail_count; response }
-  end
-
-  def no_sleep_retry(*args, **kwargs)
-    mw = Brute::Middleware::Retry.new(*args, **kwargs)
-    mw.define_singleton_method(:sleep) { |_| }
-    mw
-  end
-
-  it "returns the response on first successful call" do
-    response = MockResponse.new(content: "success")
-    app, calls = mock_inner_app(response: response)
-    middleware = Brute::Middleware::Retry.new(app)
-    result = middleware.call(build_env)
-    result.should == response
-  end
-
-  it "calls inner app exactly once on success" do
-    response = MockResponse.new(content: "success")
-    app, calls = mock_inner_app(response: response)
-    Brute::Middleware::Retry.new(app).call(build_env)
-    calls.size.should == 1
-  end
-
-  it "retries on LLM::RateLimitError and succeeds" do
-    response = MockResponse.new(content: "success")
-    app = flaky_inner_app(LLM::RateLimitError, fail_count: 2, response: response)
-    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 2)
-    env = build_env
-    result = middleware.call(env)
-    result.should == response
-  end
-
-  it "records retry_attempt in metadata after retries" do
-    response = MockResponse.new(content: "success")
-    app = flaky_inner_app(LLM::RateLimitError, fail_count: 2, response: response)
-    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 2)
-    env = build_env
-    middleware.call(env)
-    env[:metadata][:retry_attempt].should == 2
-  end
-
-  it "retries on LLM::ServerError and succeeds" do
-    response = MockResponse.new(content: "success")
-    app = flaky_inner_app(LLM::ServerError, fail_count: 1, response: response)
-    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 2)
-    result = middleware.call(build_env)
-    result.should == response
-  end
-
-  it "re-raises after exhausting all attempts" do
-    app = ->(_env) { raise LLM::RateLimitError, "rate limited" }
-    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 2)
-    lambda { middleware.call(build_env) }.should.raise(LLM::RateLimitError)
-  end
-
-  it "does not retry non-retryable errors" do
-    call_count = 0
-    app = ->(_env) { call_count += 1; raise ArgumentError, "bad input" }
-    middleware = Brute::Middleware::Retry.new(app)
-    lambda { middleware.call(build_env) }.should.raise(ArgumentError)
-  end
-
-  it "only calls inner app once for non-retryable errors" do
-    call_count = 0
-    app = ->(_env) { call_count += 1; raise ArgumentError, "bad input" }
-    middleware = Brute::Middleware::Retry.new(app)
-    begin; middleware.call(build_env); rescue ArgumentError; end
-    call_count.should == 1
-  end
-
-  it "records retry_delay in metadata" do
-    response = MockResponse.new(content: "success")
-    app = flaky_inner_app(LLM::RateLimitError, fail_count: 1, response: response)
-    middleware = no_sleep_retry(app, max_attempts: 3, base_delay: 3)
-    env = build_env
-    middleware.call(env)
-    env[:metadata][:retry_delay].should == 3
-  end
-
-  it "tracks sleep delays for exponential backoff" do
-    response = MockResponse.new(content: "success")
-    app = flaky_inner_app(LLM::RateLimitError, fail_count: 2, response: response)
-    delays = []
-    mw = Brute::Middleware::Retry.new(app, max_attempts: 3, base_delay: 2)
-    mw.define_singleton_method(:sleep) { |d| delays << d }
-    mw.call(build_env)
-    delays.should == [2, 4]
-  end
-end

data/lib/brute/middleware/session_persistence.rb

@@ -1,72 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    # Saves the conversation to disk after each LLM call.
-    #
-    # Runs POST-call: serializes env[:messages] via Session#save_messages.
-    # Failures are non-fatal — a broken session save should never crash
-    # the agent loop.
-    #
-    class SessionPersistence < Base
-      def initialize(app, session:)
-        super(app)
-        @session = session
-      end
-
-      def call(env)
-        response = @app.call(env)
-
-        begin
-          @session.save_messages(env[:messages])
-        rescue => e
-          warn "[brute] Session save failed: #{e.message}"
-        end
-
-        response
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../spec/support/mock_provider"
-  require_relative "../../../spec/support/mock_response"
-
-  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 "passes the response through unchanged" do
-    response = MockResponse.new(content: "saved response")
-    session = Struct.new(:saved) { def save_messages(m); self.saved = m; end }.new
-    inner_app = ->(_env) { response }
-    middleware = Brute::Middleware::SessionPersistence.new(inner_app, session: session)
-    result = middleware.call(build_env)
-    result.should == response
-  end
-
-  it "calls session.save_messages with env messages" do
-    response = MockResponse.new(content: "saved response")
-    session = Struct.new(:saved) { def save_messages(m); self.saved = m; end }.new
-    inner_app = ->(_env) { response }
-    middleware = Brute::Middleware::SessionPersistence.new(inner_app, session: session)
-    messages = [LLM::Message.new(:user, "hello")]
-    middleware.call(build_env(messages: messages))
-    session.saved.should == messages
-  end
-
-  it "does not propagate session save failures" do
-    response = MockResponse.new(content: "saved response")
-    session = Object.new
-    session.define_singleton_method(:save_messages) { |_| raise RuntimeError, "disk full" }
-    inner_app = ->(_env) { response }
-    middleware = Brute::Middleware::SessionPersistence.new(inner_app, session: session)
-    lambda { middleware.call(build_env) }.should.not.raise
-  end
-end