brute 1.0.1 → 2.0.1

data/lib/brute/middleware/llm_call.rb

@@ -1,128 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    # The terminal "app" in the pipeline — performs the actual LLM call.
-    #
-    # Builds a fresh LLM::Context per call from env[:messages], makes the
-    # call, extracts new messages back into env[:messages], and stashes
-    # pending functions in env[:pending_functions].
-    #
-    # When streaming, on_content fires incrementally via AgentStream.
-    # When not streaming, fires on_content post-hoc with the full text.
-    #
-    class LLMCall
-      def call(env)
-        ctx = build_context(env)
-
-        # Load existing conversation history into the ephemeral context
-        ctx.messages.concat(env[:messages])
-
-        response = ctx.talk(env[:input])
-
-        # Extract new messages appended by talk() and store them
-        new_messages = ctx.messages.to_a.drop(env[:messages].size)
-        env[:messages].concat(new_messages)
-
-        # Stash pending functions for the agent loop
-        env[:pending_functions] = ctx.functions.to_a
-
-        # Only fire on_content post-hoc when NOT streaming
-        # (streaming delivers chunks incrementally via AgentStream)
-        unless env[:streaming]
-          if (cb = env.dig(:callbacks, :on_content)) && response
-            text = safe_content(response)
-            cb.call(text) if text
-          end
-        end
-
-        response
-      end
-
-      private
-
-      def build_context(env)
-        params = {}
-        params[:tools]  = env[:tools]   if env[:tools]&.any?
-        params[:stream] = env[:stream]  if env[:stream]
-        params[:model]  = env[:model]   if env[:model]
-        LLM::Context.new(env[:provider], **params)
-      end
-
-      # Safely extract text content from an LLM response.
-      # Returns nil when the response contains only tool calls (no assistant text),
-      # which causes LLM::Contract::Completion#content to raise NoMethodError
-      # because messages.find(&:assistant?) returns nil.
-      def safe_content(response)
-        return nil unless response.respond_to?(:content)
-        response.content
-      rescue NoMethodError
-        nil
-      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 "calls the provider and returns a response" do
-    provider = MockProvider.new
-    middleware = Brute::Middleware::LLMCall.new
-    env = build_env(provider: provider, input: "hello", streaming: false)
-    response = middleware.call(env)
-    response.should.not.be.nil
-  end
-
-  it "records a call on the provider" do
-    provider = MockProvider.new
-    middleware = Brute::Middleware::LLMCall.new
-    env = build_env(provider: provider, input: "hello", streaming: false)
-    middleware.call(env)
-    provider.calls.size.should == 1
-  end
-
-  it "appends new messages to env[:messages]" do
-    provider = MockProvider.new
-    middleware = Brute::Middleware::LLMCall.new
-    env = build_env(provider: provider, input: "hello", streaming: false)
-    middleware.call(env)
-    env[:messages].should.not.be.empty
-  end
-
-  it "populates env[:pending_functions] as an Array" do
-    provider = MockProvider.new
-    middleware = Brute::Middleware::LLMCall.new
-    env = build_env(provider: provider, input: "hello", streaming: false)
-    middleware.call(env)
-    env[:pending_functions].should.be.kind_of(Array)
-  end
-
-  it "does not fire on_content callback when streaming" do
-    provider = MockProvider.new
-    middleware = Brute::Middleware::LLMCall.new
-    called = false
-    env = build_env(provider: provider, input: "hi", streaming: true, callbacks: { on_content: ->(_) { called = true } })
-    middleware.call(env)
-    called.should.be.false
-  end
-
-  it "preserves existing messages across calls" do
-    provider = MockProvider.new
-    middleware = Brute::Middleware::LLMCall.new
-    existing = LLM::Message.new(:user, "previous")
-    env = build_env(provider: provider, input: "hello", streaming: false, messages: [existing])
-    middleware.call(env)
-    env[:messages].first.should == existing
-  end
-end

data/lib/brute/middleware/message_tracking.rb

@@ -1,339 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    # Records every LLM exchange into a MessageStore in the OpenCode
-    # {info, parts} format so sessions can be viewed later.
-    #
-    # Lifecycle per pipeline call:
-    #
-    #   1. PRE-CALL  — if this is the first call of a turn (env[:tool_results]
-    #      is nil), record the user message.
-    #   2. POST-CALL — record the assistant message: text content as a "text"
-    #      part, each tool call as a "tool" part in "running" state.
-    #   3. When the pipeline is called again with tool results, update the
-    #      corresponding tool parts to "completed" (or "error").
-    #
-    # The middleware also stores itself in env[:message_tracking] so the
-    # agent loop can access the current assistant message ID for callbacks.
-    #
-    class MessageTracking < Base
-      attr_reader :store
-
-      def initialize(app, store:)
-        super(app)
-        @store = store
-        @current_user_id = nil
-        @current_assistant_id = nil
-      end
-
-      def call(env)
-        env[:message_tracking] = self
-
-        # ── Pre-call: record user message or update tool results ──
-        if env[:tool_results].nil?
-          # New turn — record the user message
-          record_user_message(env)
-        else
-          # Tool results coming back — complete the tool parts
-          complete_tool_parts(env)
-        end
-
-        # ── LLM call ──
-        response = @app.call(env)
-
-        # ── Post-call: record assistant message ──
-        record_assistant_message(env, response)
-
-        response
-      end
-
-      # The current assistant message ID (used by external callbacks).
-      def current_assistant_id
-        @current_assistant_id
-      end
-
-      private
-
-      # ── User message ───────────────────────────────────────────────
-
-      def record_user_message(env)
-        text = extract_user_text(env)
-        return unless text
-
-        @current_user_id = @store.append_user(text: text)
-      end
-
-      def extract_user_text(env)
-        input = env[:input]
-        case input
-        when String
-          input
-        when Array
-          # llm.rb prompt format: array of message hashes
-          user_msg = input.reverse_each.find { |m| m.respond_to?(:role) && m.role.to_s == "user" }
-          user_msg&.content.to_s if user_msg
-        else
-          # Could be a prompt object — try to extract user content
-          if input.respond_to?(:messages)
-            msgs = input.messages.to_a
-            user_msg = msgs.reverse_each.find { |m| m.role.to_s == "user" }
-            user_msg&.content.to_s if user_msg
-          end
-        end
-      end
-
-      # ── Assistant message ──────────────────────────────────────────
-
-      def record_assistant_message(env, response)
-        provider_name = env[:provider]&.class&.name&.split("::")&.last&.downcase
-        model_name = resolve_model_name(env)
-
-        @current_assistant_id = @store.append_assistant(
-          parent_id: @current_user_id,
-          model_id: model_name,
-          provider_id: provider_name,
-        )
-
-        # Text content
-        text = safe_content(response)
-        @store.add_text_part(message_id: @current_assistant_id, text: text) if text && !text.empty?
-
-        # Tool calls
-        record_tool_calls(env)
-
-        # Token usage
-        tokens = extract_tokens(env, response)
-        @store.complete_assistant(message_id: @current_assistant_id, tokens: tokens) if tokens
-
-        # Step finish
-        @store.add_step_finish(message_id: @current_assistant_id, tokens: tokens)
-      end
-
-      def record_tool_calls(env)
-        functions = env[:pending_functions]
-        return if functions.nil? || functions.empty?
-
-        functions.each do |fn|
-          @store.add_tool_part(
-            message_id: @current_assistant_id,
-            tool: fn.name,
-            call_id: fn.id,
-            input: fn.arguments,
-          )
-        end
-      end
-
-      # ── Tool results ───────────────────────────────────────────────
-
-      def complete_tool_parts(env)
-        return unless @current_assistant_id
-
-        results = env[:tool_results]
-        return unless results.is_a?(Array)
-
-        results.each do |name, value|
-          # Find the tool part by name (tool results come as [name, value] pairs)
-          msg = @store.message(@current_assistant_id)
-          next unless msg
-
-          # Match by tool name — find the first running tool part with this name
-          part = msg[:parts]&.find do |p|
-            p[:type] == "tool" && p[:tool] == name && p.dig(:state, :status) == "running"
-          end
-          next unless part
-
-          call_id = part[:callID]
-          if value.is_a?(Hash) && value[:error]
-            @store.error_tool_part(
-              message_id: @current_assistant_id,
-              call_id: call_id,
-              error: value[:error],
-            )
-          else
-            output = value.is_a?(String) ? value : value.to_s
-            @store.complete_tool_part(
-              message_id: @current_assistant_id,
-              call_id: call_id,
-              output: output,
-            )
-          end
-        end
-      end
-
-      # ── Helpers ────────────────────────────────────────────────────
-
-      # Resolve the actual model used for the request.
-      # Prefers env[:model] (set by AgentTurn) and falls back to the
-      # provider's default_model.
-      def resolve_model_name(env)
-        model = env[:model]
-        return model.to_s if model
-
-        # Fall back to provider default
-        env[:provider]&.respond_to?(:default_model) ? env[:provider].default_model.to_s : nil
-      end
-
-      def safe_content(response)
-        return nil unless response.respond_to?(:content)
-        response.content
-      rescue NoMethodError
-        nil
-      end
-
-      def extract_tokens(env, response)
-        # Prefer the metadata accumulated by TokenTracking middleware
-        meta_tokens = env.dig(:metadata, :tokens, :last_call)
-        if meta_tokens
-          {
-            input: meta_tokens[:input] || 0,
-            output: meta_tokens[:output] || 0,
-            reasoning: 0,
-            cache: { read: 0, write: 0 },
-          }
-        elsif response.respond_to?(:usage) && (usage = response.usage)
-          {
-            input: usage.input_tokens.to_i,
-            output: usage.output_tokens.to_i,
-            reasoning: usage.reasoning_tokens.to_i,
-            cache: { read: 0, write: 0 },
-          }
-        end
-      end
-    end
-  end
-end
-
-test do
-  require_relative "../../../spec/support/mock_provider"
-  require_relative "../../../spec/support/mock_response"
-  require "tmpdir"
-  require "fileutils"
-
-  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 with_tracking
-    tmpdir = Dir.mktmpdir("brute_test_")
-    store = Brute::Store::MessageStore.new(session_id: "test-session", dir: tmpdir)
-    response = MockResponse.new(content: "Hello from the LLM")
-    inner_app = ->(_env) { response }
-    middleware = Brute::Middleware::MessageTracking.new(inner_app, store: store)
-    yield middleware, store, response
-  ensure
-    FileUtils.rm_rf(tmpdir)
-  end
-
-  it "records a user message on first call of a turn" do
-    with_tracking do |mw, store, _|
-      mw.call(build_env(input: "What is Ruby?", tool_results: nil))
-      user_msg = store.messages.find { |m| m[:info][:role] == "user" }
-      user_msg[:parts][0][:text].should == "What is Ruby?"
-    end
-  end
-
-  it "records only one user message per turn" do
-    with_tracking do |mw, store, _|
-      env = build_env(input: "Hello", tool_results: nil)
-      mw.call(env)
-      env[:tool_results] = [["read", "contents"]]
-      mw.call(env)
-      store.messages.select { |m| m[:info][:role] == "user" }.size.should == 1
-    end
-  end
-
-  it "records an assistant message after LLM call" do
-    with_tracking do |mw, store, _|
-      mw.call(build_env(input: "Hello", tool_results: nil))
-      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
-      asst.should.not.be.nil
-    end
-  end
-
-  it "captures text content as a text part" do
-    with_tracking do |mw, store, _|
-      mw.call(build_env(input: "Hello", tool_results: nil))
-      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
-      text_parts = asst[:parts].select { |p| p[:type] == "text" }
-      text_parts[0][:text].should == "Hello from the LLM"
-    end
-  end
-
-  it "captures token usage from response" do
-    with_tracking do |mw, store, _|
-      mw.call(build_env(input: "Hello", tool_results: nil))
-      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
-      asst[:info][:tokens][:input].should == 100
-    end
-  end
-
-  it "records tool calls as tool parts in running state" do
-    with_tracking do |mw, store, _|
-      fn = Struct.new(:id, :name, :arguments, keyword_init: true).new(id: "call_001", name: "read", arguments: { file_path: "/test" })
-      mw.call(build_env(input: "Read the file", tool_results: nil, pending_functions: [fn]))
-      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
-      tool_parts = asst[:parts].select { |p| p[:type] == "tool" }
-      tool_parts[0][:state][:status].should == "running"
-    end
-  end
-
-  it "updates tool parts when results arrive" do
-    with_tracking do |mw, store, _|
-      fn = Struct.new(:id, :name, :arguments, keyword_init: true).new(id: "call_001", name: "read", arguments: { file_path: "/test" })
-      env = build_env(input: "Read the file", tool_results: nil, pending_functions: [fn])
-      mw.call(env)
-      env[:pending_functions] = []
-      env[:tool_results] = [["read", "file contents here"]]
-      mw.call(env)
-      first_asst = store.messages.find { |m| m[:info][:role] == "assistant" }
-      tool_part = first_asst[:parts].find { |p| p[:type] == "tool" }
-      tool_part[:state][:status].should == "completed"
-    end
-  end
-
-  it "records provider default_model when no override" do
-    with_tracking do |mw, store, _|
-      mw.call(build_env(input: "Hello", tool_results: nil))
-      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
-      asst[:info][:modelID].should == "mock-model"
-    end
-  end
-
-  it "records overridden model when env[:model] is set" do
-    with_tracking do |mw, store, _|
-      mw.call(build_env(input: "Hello", tool_results: nil, model: "custom-haiku"))
-      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
-      asst[:info][:modelID].should == "custom-haiku"
-    end
-  end
-
-  it "stores itself in env[:message_tracking]" do
-    with_tracking do |mw, _, _|
-      env = build_env(input: "Hello", tool_results: nil)
-      mw.call(env)
-      env[:message_tracking].should == mw
-    end
-  end
-
-  it "returns the inner app response unchanged" do
-    with_tracking do |mw, _, response|
-      result = mw.call(build_env(input: "Hello", tool_results: nil))
-      result.should == response
-    end
-  end
-
-  it "adds a step-finish part to assistant messages" do
-    with_tracking do |mw, store, _|
-      mw.call(build_env(input: "Hello", tool_results: nil))
-      asst = store.messages.find { |m| m[:info][:role] == "assistant" }
-      step_finish = asst[:parts].find { |p| p[:type] == "step-finish" }
-      step_finish[:reason].should == "stop"
-    end
-  end
-end

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

@@ -1,105 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    module OTel
-      # Outermost OTel middleware. Creates a span per LLM pipeline call
-      # and passes it through env[:span] for inner OTel middlewares to
-      # decorate with events and attributes.
-      #
-      # When opentelemetry-sdk is not loaded, this is a pure pass-through.
-      #
-      # Pipeline position: outermost (wraps everything including retries).
-      #
-      #   use Brute::Middleware::OTel::Span
-      #   use Brute::Middleware::OTel::ToolResults
-      #   use Brute::Middleware::OTel::ToolCalls
-      #   use Brute::Middleware::OTel::TokenUsage
-      #   # ... existing middleware ...
-      #   run Brute::Middleware::LLMCall.new
-      #
-      class Span < Base
-        def call(env)
-          return @app.call(env) unless defined?(::OpenTelemetry::SDK)
-
-          provider_name = provider_type(env[:provider])
-          model = env[:model] || (env[:provider].default_model rescue nil)
-          span_name = model ? "llm.call #{model}" : "llm.call"
-
-          attributes = {
-            "brute.provider" => provider_name,
-            "brute.streaming" => !!env[:streaming],
-            "brute.context_messages" => env[:messages].size,
-          }
-          attributes["brute.model"] = model.to_s if model
-          attributes["brute.session_id"] = env[:metadata][:session_id].to_s if env.dig(:metadata, :session_id)
-
-          tracer.in_span(span_name, attributes: attributes, kind: :internal) do |span|
-            env[:span] = span
-            response = @app.call(env)
-
-            # Record response model if it differs from request model
-            resp_model = begin; response.model; rescue; nil; end
-            span.set_attribute("brute.response_model", resp_model.to_s) if resp_model && resp_model != model
-
-            response
-          rescue ::StandardError => e
-            span.record_exception(e)
-            span.status = ::OpenTelemetry::Trace::Status.error(e.message)
-            raise
-          ensure
-            env.delete(:span)
-          end
-        end
-
-        private
-
-        def tracer
-          @tracer ||= ::OpenTelemetry.tracer_provider.tracer("brute", Brute::VERSION)
-        end
-
-        def provider_type(provider)
-          name = provider.class.name.to_s.downcase
-          if name.include?("anthropic") then "anthropic"
-          elsif name.include?("openai") then "openai"
-          elsif name.include?("google") || name.include?("gemini") then "google"
-          elsif name.include?("deepseek") then "deepseek"
-          elsif name.include?("ollama") then "ollama"
-          elsif name.include?("xai") then "xai"
-          else "unknown"
-          end
-        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 through when OpenTelemetry::SDK is not defined" do
-    response = MockResponse.new(content: "hello from LLM")
-    middleware = Brute::Middleware::OTel::Span.new(->(_env) { response })
-    env = build_env
-    result = middleware.call(env)
-    result.should == response
-  end
-
-  it "env[:span] is nil when OTel is not defined" do
-    response = MockResponse.new(content: "hello from LLM")
-    middleware = Brute::Middleware::OTel::Span.new(->(_env) { response })
-    env = build_env
-    middleware.call(env)
-    env[:span].should.be.nil
-  end
-end

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

@@ -1,68 +0,0 @@
-# frozen_string_literal: true
-
-require "bundler/setup"
-require "brute"
-
-module Brute
-  module Middleware
-    module OTel
-      # Records token usage from the LLM response as span attributes.
-      #
-      # Runs POST-call: reads token counts from the response usage object
-      # and sets them as attributes on the span.
-      #
-      class TokenUsage < Base
-        def call(env)
-          response = @app.call(env)
-
-          span = env[:span]
-          if span && response.respond_to?(:usage) && (usage = response.usage)
-            span.set_attribute("gen_ai.usage.input_tokens", usage.input_tokens.to_i)
-            span.set_attribute("gen_ai.usage.output_tokens", usage.output_tokens.to_i)
-            span.set_attribute("gen_ai.usage.total_tokens", usage.total_tokens.to_i)
-
-            reasoning = usage.reasoning_tokens.to_i
-            span.set_attribute("gen_ai.usage.reasoning_tokens", reasoning) if reasoning > 0
-          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
-
-  def make_response
-    MockResponse.new(content: "hello",
-      usage: LLM::Usage.new(input_tokens: 100, output_tokens: 50, reasoning_tokens: 10, total_tokens: 160))
-  end
-
-  it "passes the response through unchanged" do
-    response = make_response
-    middleware = Brute::Middleware::OTel::TokenUsage.new(->(_env) { response })
-    result = middleware.call(build_env)
-    result.should == response
-  end
-
-  it "passes through without error when span is nil" do
-    response = make_response
-    middleware = Brute::Middleware::OTel::TokenUsage.new(->(_env) { response })
-    lambda { middleware.call(build_env) }.should.not.raise
-  end
-
-  it "handles a response without usage gracefully" do
-    no_usage = Object.new
-    middleware = Brute::Middleware::OTel::TokenUsage.new(->(_env) { no_usage })
-    lambda { middleware.call(build_env) }.should.not.raise
-  end
-end