brute 5.0.5 → 6.0.0

data/lib/brute/token_counter/tiktoken.rb

@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "brute"
+
+module Brute
+  module TokenCounter
+    # Tokens from OpenAI's byte-pair encoder, through the tiktoken_ruby gem.
+    #
+    #   gem "tiktoken_ruby"
+    #   use Brute::Middleware::DefaultCompactionPipeline,
+    #     window:        200_000,
+    #     token_counter: Brute::TokenCounter::Tiktoken.new
+    #
+    # Brute depends on no LLM library, and this is no exception: the gem is
+    # required the first time the counter is asked for a number, so an agent
+    # that never installs it never pays for it and never hears about it.
+    #
+    # It is exact about the text and still approximate about the request --
+    # `per_message` stands in for the chat-template framing, which is the
+    # provider's and not in any encoder.
+    class Tiktoken
+      ENCODING = "o200k_base"
+
+      def initialize(encoding: ENCODING, per_message: 4, encoder: nil)
+        @encoding = encoding
+        @per_message = per_message
+        @encoder = encoder
+      end
+
+      # Load the encoder, downloading its vocabulary if it is not cached yet.
+      def warm_up
+        @encoder ||= load_encoder
+      end
+
+      def count(messages, tools: nil)
+        warm_up
+        text = Rendering.conversation(messages) + Rendering.tools(tools)
+
+        @encoder.encode(text).length + (Array(messages).length * @per_message)
+      end
+
+      private
+
+        def load_encoder
+          begin
+            require "tiktoken_ruby"
+          rescue LoadError
+            raise LoadError, "#{self.class} needs the tiktoken_ruby gem: add `gem \"tiktoken_ruby\"` to your Gemfile"
+          end
+
+          ::Tiktoken.get_encoding(@encoding)
+        end
+    end
+  end
+end
+
+__END__
+
+describe "brute/token_counter/tiktoken" do
+  it "encodes the same rendering the approximate counter measures, and says what to install when it cannot" do
+    encoded = []
+    encoder = Object.new
+    encoder.define_singleton_method(:encode) { |text| encoded << text; text.split(/\b/) }
+
+    log = Brute.log
+    log.user("what changed?")
+
+    counter = Brute::TokenCounter::Tiktoken.new(encoder: encoder, per_message: 4)
+    counter.count(log).should == encoder.encode("user: what changed?").length + 4
+    encoded.last.should == "user: what changed?"
+
+    # The vocabulary is loaded once, on the first count rather than at boot.
+    counter.warm_up.should.equal? encoder
+
+    missing = Brute::TokenCounter::Tiktoken.new(encoding: "nothing-real")
+    missing.define_singleton_method(:require) { |_name| raise LoadError }
+    lambda { missing.count(log) }.should.raise(LoadError).message.should.include "tiktoken_ruby"
+  end
+end

data/lib/brute/token_counter.rb

@@ -0,0 +1,150 @@
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "brute"
+require "json"
+
+module Brute
+  # How big a conversation is, in tokens.
+  #
+  # A counter is anything answering `count(messages, tools: nil)`. The tools
+  # are part of the question because their schemas ride in every request
+  # alongside the messages -- an agent carrying a dozen of them is spending
+  # context on JSON Schema before anyone has said a word.
+  #
+  #   counter = Brute::TokenCounter::Approximate.new
+  #   counter.count(env[:messages], tools: env[:tools])
+  #
+  # What a turn is *currently* costing is a different question, and
+  # `estimate` answers it: the provider counted the conversation exactly when
+  # it answered, so that number is trusted and only what has been appended
+  # since is counted locally.
+  #
+  #   Brute::TokenCounter.estimate(env)
+  #
+  module TokenCounter
+    def self.default = Approximate.new
+
+    # What the whole turn costs right now.
+    #
+    # :counter defaults to the one the turn already decided on
+    # (env[:token_counter]), :tools to the ones the pipeline advertises
+    # (env[:tools]).
+    #
+    # The reported total already covers the system prompt, the tool schemas
+    # and the provider's own chat-template overhead, so the warm path adds
+    # only the messages that landed after the reply it describes -- and never
+    # the schemas, which are already inside it. Anything else double counts.
+    def self.estimate(env, counter: nil, tools: nil)
+      counter ||= env[:token_counter] || default
+      if tools.nil?
+        tools = env[:tools]
+      end
+      messages = env[:messages] || []
+      reported = env.dig(:metadata, :last_llm_usage)&.total.to_i
+      answered = messages.rindex { |message| message.role == :assistant }
+
+      # Nothing sent yet, or a conversation the reported number no longer
+      # describes -- a compaction that left no reply behind, say. Count it all.
+      if reported.zero? || answered.nil?
+        counter.count(messages, tools: tools)
+      else
+        reported + counter.count(messages[(answered + 1)..])
+      end
+    end
+
+    # The text a counter measures.
+    #
+    # One rendering, so two counters cannot disagree about what the
+    # conversation even is -- and it is the same text a summariser reads,
+    # which is why it says who spoke and what was called rather than being
+    # the shortest thing that could be measured.
+    module Rendering
+      def self.conversation(messages) = Array(messages).map { |message| message(message) }.join("\n")
+
+      def self.message(message)
+        parts = ["#{message.role}:", message.content.to_s]
+
+        message.tool_calls&.each do |call|
+          parts << "#{call.name}(#{JSON.generate(call.arguments)}) -> #{call.id}"
+        end
+
+        unless message.tool_call_id.nil?
+          parts << "(answering #{message.tool_call_id})"
+        end
+
+        parts.reject(&:empty?).join(" ")
+      end
+
+      # The schemas as the provider is given them, which is what they cost.
+      def self.tools(tools)
+        wrapped = Brute::Tools::Adapter.wrap_all(tools || [])
+
+        if wrapped.empty?
+          ""
+        else
+          JSON.generate(wrapped.values.map(&:to_h))
+        end
+      end
+    end
+  end
+end
+
+__END__
+
+describe "brute/token_counter" do
+  def usage(total) = Brute::UsageDetection::Usage.new(total: total)
+
+  it "measures a conversation and its schemas, and trusts the provider for what it already counted" do
+    log = Brute.log
+    log.system("you are a helpful agent")
+    log.user("what changed?")
+
+    counter = Brute::TokenCounter.default
+    counter.should.be.kind_of Brute::TokenCounter::Approximate
+
+    text = Brute::TokenCounter::Rendering.conversation(log)
+    text.should == "system: you are a helpful agent\nuser: what changed?"
+
+    # A tool call and the result answering it are rendered too -- both cost
+    # tokens, and neither is in the message's content.
+    log << Brute::Message.new(
+      role: :assistant,
+      content: "",
+      tool_calls: [{ id: "tc1", name: "search", arguments: { "query" => "fed" } }]
+    )
+    log.tool("found nothing", tool_call_id: "tc1")
+    Brute::TokenCounter::Rendering.conversation(log).should.include 'search({"query":"fed"}) -> tc1'
+    Brute::TokenCounter::Rendering.conversation(log).should.include "(answering tc1)"
+
+    # Schemas ride in every request, so they are part of the question.
+    tool = Brute::Turn::ToolPipeline.new(name: "search", description: "search the web") do
+      run ->(env) { env[:result] = "" }
+    end
+    Brute::TokenCounter::Rendering.tools(nil).should == ""
+    Brute::TokenCounter::Rendering.tools([tool]).should.include "search the web"
+    counter.count(log, tools: [tool]).should > counter.count(log)
+
+    # Nothing reported yet: everything is counted, schemas included.
+    cold = { messages: log, tools: [tool], metadata: {} }
+    Brute::TokenCounter.estimate(cold).should == counter.count(log, tools: [tool])
+
+    # Reported: the provider's number, plus only what landed after the reply
+    # it describes. The schemas are already inside it.
+    log.assistant("nothing changed")
+    log.user("are you sure?")
+    warm = { messages: log, tools: [tool], metadata: { last_llm_usage: usage(9_000) } }
+    Brute::TokenCounter.estimate(warm).should == 9_000 + counter.count([log.last])
+
+    # A conversation the reported number no longer describes -- a compaction
+    # that left no reply behind -- is counted from scratch rather than added
+    # to a total for a conversation that is gone.
+    stale = { messages: Brute.log.tap { |l| l.user("only me") }, metadata: { last_llm_usage: usage(9_000) } }
+    Brute::TokenCounter.estimate(stale).should == counter.count(stale[:messages])
+
+    # The counter the turn already decided on is the one that gets used.
+    given = { messages: log, metadata: {}, token_counter: Object.new }
+    given[:token_counter].define_singleton_method(:count) { |_messages, tools: nil| 7 }
+    Brute::TokenCounter.estimate(given).should == 7
+  end
+end

data/lib/brute/tool.rb

@@ -30,9 +30,11 @@ module Brute
   class Tool
     class << self
       def description(text = nil)
-        return @description unless text
-
-        @description = text
+        if text
+          @description = text
+        else
+          @description
+        end
       end
 
       # Declare one parameter: param :key, type:, desc:, required:
@@ -46,9 +48,11 @@ module Brute
 
       # Raw JSON-schema override for complex argument shapes.
       def params(schema = nil)
-        return @params_schema unless schema
-
-        @params_schema = schema.deep_symbolize_keys
+        if schema
+          @params_schema = schema.deep_symbolize_keys
+        else
+          @params_schema
+        end
       end
     end
 

data/lib/brute/tools/adapter.rb

@@ -39,29 +39,32 @@ module Brute
 
       # Wrap a single tool of any supported shape. Idempotent.
       def self.wrap(tool)
-        return tool if tool.is_a?(Adapter)
-
-        tool = tool.new if tool.is_a?(Class)
-
-        case tool
-        when Hash                then from_hash(tool)
-        when ::Brute::Tool       then from_brute_tool(tool)
-        when Brute::Tools::SubAgent then new(
-          name:        tool.name,
-          description: tool.description,
-          params:      tool.params,
-          handler:     ->(**args) { tool.execute(args) },
-          original:    tool,
-        )
-        when Brute::Turn::ToolPipeline then new(
-          name:        tool.name,
-          description: tool.description,
-          params:      tool.params,
-          handler:     ->(**args) { tool.call(**args) },
-          original:    tool,
-        )
+        if tool.is_a?(Adapter)
+          tool
         else
-          from_duck_type(tool)
+          if tool.is_a?(Class)
+            tool = tool.new
+          end
+          case tool
+          when Hash                then from_hash(tool)
+          when ::Brute::Tool       then from_brute_tool(tool)
+          when Brute::Tools::SubAgent then new(
+            name:        tool.name,
+            description: tool.description,
+            params:      tool.params,
+            handler:     ->(**args) { tool.execute(args) },
+            original:    tool,
+          )
+          when Brute::Turn::ToolPipeline then new(
+            name:        tool.name,
+            description: tool.description,
+            params:      tool.params,
+            handler:     ->(**args) { tool.call(**args) },
+            original:    tool,
+          )
+          else
+            from_duck_type(tool)
+          end
         end
       end
 
@@ -78,7 +81,9 @@ module Brute
       def self.from_hash(definition)
         definition = definition.transform_keys(&:to_sym)
         handler = definition.fetch(:execute) { definition[:handler] }
-        raise ArgumentError, "inline tool needs an :execute proc" unless handler.respond_to?(:call)
+        unless handler.respond_to?(:call)
+          raise ArgumentError, "inline tool needs an :execute proc"
+        end
 
         new(
           name:        definition.fetch(:name).to_s,
@@ -108,7 +113,11 @@ module Brute
           raise ArgumentError, "don't know how to adapt #{tool.inspect} into a tool"
         end
 
-        entry = tool.respond_to?(:execute) ? tool.method(:execute) : tool.method(:call)
+        if tool.respond_to?(:execute)
+          entry = tool.method(:execute)
+        else
+          entry = tool.method(:call)
+        end
         new(
           name:        tool.name.to_s,
           description: tool.respond_to?(:description) ? tool.description : "",
@@ -118,6 +127,10 @@ module Brute
         )
       end
 
+      # The tool object this adapter wraps (Brute::Tool, Brute::Turn::ToolPipeline,
+      # SubAgent, Hash definition, ...).
+      attr_reader :original
+
       def initialize(name:, description:, params:, handler:, schema: nil, original: nil)
         @name        = name
         @description = description
@@ -127,10 +140,6 @@ module Brute
         @original    = original
       end
 
-      # The tool object this adapter wraps (Brute::Tool, Brute::Turn::ToolPipeline,
-      # SubAgent, Hash definition, ...).
-      attr_reader :original
-
       # Execute the tool. Accepts string- or symbol-keyed argument hashes,
       # as delivered by LLM providers.
       def call(arguments = {})
@@ -141,27 +150,28 @@ module Brute
       # Library-neutral tool definition (JSON-Schema-ish). The inline `run`
       # proc reshapes this into whatever its LLM library expects.
       def to_h
-        return { name: @name, description: @description, parameters: @schema.deep_symbolize_keys } if @schema
-
-        properties = @params.transform_values do |opts|
+        if @schema
+          { name: @name, description: @description, parameters: @schema.deep_symbolize_keys }
+        else
+          properties = @params.transform_values do |opts|
+            {
+              type:        opts[:type] || "string",
+              description: opts[:desc] || opts[:description],
+              items:       opts[:items],
+              enum:        opts[:enum],
+            }.compact
+          end
+          required = @params.select { |_k, opts| opts[:required] }.keys
           {
-            type:        opts[:type] || "string",
-            description: opts[:desc] || opts[:description],
-            items:       opts[:items],
-            enum:        opts[:enum],
-          }.compact
+            name:        @name,
+            description: @description,
+            parameters:  {
+              type:       "object",
+              properties: properties,
+              required:   required.map(&:to_s),
+            },
+          }
         end
-        required = @params.select { |_k, opts| opts[:required] }.keys
-
-        {
-          name:        @name,
-          description: @description,
-          parameters: {
-            type:       "object",
-            properties: properties,
-            required:   required.map(&:to_s),
-          },
-        }
       end
     end
   end

data/lib/brute/tools/fs/snapshot_store.rb

@@ -19,7 +19,11 @@ module Brute
           # If the file doesn't exist yet, records +:did_not_exist+.
           def save(path)
             key = File.expand_path(path)
-            content = File.exist?(key) ? File.read(key) : :did_not_exist
+            if File.exist?(key)
+              content = File.read(key)
+            else
+              content = :did_not_exist
+            end
             @mutex.synchronize { @snapshots[key].push(content) }
           end
 

data/lib/brute/tools/fs_patch.rb

@@ -20,22 +20,30 @@ module Brute
       def execute(file_path:, old_string:, new_string:, replace_all: false)
         path = File.expand_path(file_path)
         Brute::Tools::FS::FileMutationQueue.serialize(path) do
-          raise "File not found: #{path}" unless File.exist?(path)
+          unless File.exist?(path)
+            raise "File not found: #{path}"
+          end
 
           original = File.read(path)
-          raise "old_string not found in #{path}" unless original.include?(old_string)
+          unless original.include?(old_string)
+            raise "old_string not found in #{path}"
+          end
 
           Brute::Tools::FS::SnapshotStore.save(path)
 
-          updated = if replace_all
-                      original.gsub(old_string, new_string)
-                    else
-                      original.sub(old_string, new_string)
-                    end
+          if replace_all
+            updated = original.gsub(old_string, new_string)
+          else
+            updated = original.sub(old_string, new_string)
+          end
 
           File.write(path, updated)
           diff = Brute::Diff.unified(original, updated)
-          count = replace_all ? original.scan(old_string).size : 1
+          if replace_all
+            count = original.scan(old_string).size
+          else
+            count = 1
+          end
           { success: true, file_path: path, replacements: count, diff: diff }
         end
       end

data/lib/brute/tools/fs_read.rb

@@ -47,101 +47,128 @@ module Brute
         path = File.expand_path(file_path)
 
         # Directory listing
-        return list_directory(path) if File.directory?(path)
-
-        # File-not-found suggestions
-        unless File.exist?(path)
-          suggestions = find_similar(path)
-          msg = "File not found: #{path}"
-          msg += ". Did you mean: #{suggestions.join(', ')}?" if suggestions.any?
-          raise msg
-        end
-
-        raise "Not a file: #{path}" unless File.file?(path)
-
-        # Binary file detection
-        ext = File.extname(path).downcase
-        raise "Cannot read binary file: #{path}" if BINARY_EXTENSIONS.include?(ext)
-
-        sample = File.binread(path, 4096) || ""
-        raise "Cannot read binary file: #{path}" if sample.include?("\x00")
-
-        lines = File.readlines(path)
-        total = lines.size
-        first = start_line ? [start_line - 1, 0].max : 0
+        if File.directory?(path)
+          list_directory(path)
+        else
+          # File-not-found suggestions
+          unless File.exist?(path)
+            suggestions = find_similar(path)
+            msg = "File not found: #{path}"
+            if suggestions.any?
+              msg += ". Did you mean: #{suggestions.join(', ')}?"
+            end
+            raise msg
+          end
 
-        # Apply default line cap when no explicit range given
-        default_last = end_line ? [end_line - 1, total - 1].min : [first + DEFAULT_LINE_CAP - 1, total - 1].min
-        last = default_last
+          unless File.file?(path)
+            raise "Not a file: #{path}"
+          end
 
-        selected = lines[first..last] || []
+          # Binary file detection
+          ext = File.extname(path).downcase
+          if BINARY_EXTENSIONS.include?(ext)
+            raise "Cannot read binary file: #{path}"
+          end
+          sample = File.binread(path, 4096) || ""
+          if sample.include?("\x00")
+            raise "Cannot read binary file: #{path}"
+          end
 
-        # Per-line truncation + byte cap
-        numbered = []
-        bytes = 0
-        selected.each_with_index do |line, i|
-          truncated_line = Brute::Truncation.truncate_line(line, max: MAX_LINE_LENGTH)
-          numbered_line = "#{first + i + 1}\t#{truncated_line}"
-          break if bytes + numbered_line.bytesize > MAX_BYTES
-          numbered << numbered_line
-          bytes += numbered_line.bytesize
-        end
+          lines = File.readlines(path)
+          total = lines.size
+          if start_line
+            first = [start_line - 1, 0].max
+          else
+            first = 0
+          end
 
-        actual_last = first + numbered.size - 1
-        content = numbered.join
-        truncated = (actual_last < total - 1) && end_line.nil?
+          # Apply default line cap when no explicit range given
+          if end_line
+            default_last = [end_line - 1, total - 1].min
+          else
+            default_last = [first + DEFAULT_LINE_CAP - 1, total - 1].min
+          end
+          last = default_last
+          selected = lines[first..last] || []
+
+          # Per-line truncation + byte cap
+          numbered = []
+          bytes = 0
+          selected.each_with_index do |line, i|
+            truncated_line = Brute::Truncation.truncate_line(line, max: MAX_LINE_LENGTH)
+            numbered_line = "#{first + i + 1}\t#{truncated_line}"
+            if bytes + numbered_line.bytesize > MAX_BYTES
+              break
+            end
+            numbered << numbered_line
+            bytes += numbered_line.bytesize
+          end
 
-        if truncated
-          content + "\n(Showing lines #{first + 1}-#{actual_last + 1} of #{total}. Use start_line=#{actual_last + 2} to continue.)"
-        else
-          content
+          actual_last = first + numbered.size - 1
+          content = numbered.join
+          truncated = (actual_last < total - 1) && end_line.nil?
+          if truncated
+            content + "\n(Showing lines #{first + 1}-#{actual_last + 1} of #{total}. Use start_line=#{actual_last + 2} to continue.)"
+          else
+            content
+          end
         end
       end
 
       private
 
-      def list_directory(path)
-        entries = Dir.entries(path).reject { |e| e.start_with?(".") }.sort
-        total = entries.size
-        capped = entries.first(DEFAULT_LINE_CAP)
-        result = capped.map do |entry|
-          full = File.join(path, entry)
-          type = File.directory?(full) ? "dir" : "file"
-          "#{entry} (#{type})"
-        end.join("\n")
-
-        if total > DEFAULT_LINE_CAP
-          result += "\n(Showing #{DEFAULT_LINE_CAP} of #{total} entries)"
-        end
-        result
-      end
+        def list_directory(path)
+          entries = Dir.entries(path).reject { |e| e.start_with?(".") }.sort
+          total = entries.size
+          capped = entries.first(DEFAULT_LINE_CAP)
+          result = capped.map do |entry|
+            full = File.join(path, entry)
+            if File.directory?(full)
+              type = "dir"
+          else
+            type = "file"
+          end
+            "#{entry} (#{type})"
+          end.join("\n")
 
-      def find_similar(path)
-        dir = File.dirname(path)
-        target = File.basename(path)
-        return [] unless File.directory?(dir)
+          if total > DEFAULT_LINE_CAP
+            result += "\n(Showing #{DEFAULT_LINE_CAP} of #{total} entries)"
+          end
+          result
+        end
 
-        entries = Dir.entries(dir).reject { |e| e.start_with?(".") }
-        entries.select { |e| levenshtein(e.downcase, target.downcase) <= 3 }
-               .sort_by { |e| levenshtein(e.downcase, target.downcase) }
-               .first(3)
-      end
+        def find_similar(path)
+          dir = File.dirname(path)
+          target = File.basename(path)
+          if File.directory?(dir)
+            entries = Dir.entries(dir).reject { |e| e.start_with?(".") }
+            entries.select { |e| levenshtein(e.downcase, target.downcase) <= 3 }
+                   .sort_by { |e| levenshtein(e.downcase, target.downcase) }
+                   .first(3)
+          else
+            []
+          end
+        end
 
-      def levenshtein(a, b)
-        m, n = a.length, b.length
-        d = Array.new(m + 1) { |i| i }
-        (1..n).each do |j|
-          prev = d[0]
-          d[0] = j
-          (1..m).each do |i|
-            cost = a[i - 1] == b[j - 1] ? 0 : 1
-            temp = d[i]
-            d[i] = [d[i] + 1, d[i - 1] + 1, prev + cost].min
-            prev = temp
+        def levenshtein(a, b)
+          m, n = a.length, b.length
+          d = Array.new(m + 1) { |i| i }
+          (1..n).each do |j|
+            prev = d[0]
+            d[0] = j
+            (1..m).each do |i|
+              if a[i - 1] == b[j - 1]
+                cost = 0
+              else
+                cost = 1
+              end
+              temp = d[i]
+              d[i] = [d[i] + 1, d[i - 1] + 1, prev + cost].min
+              prev = temp
+            end
           end
+          d[m]
         end
-        d[m]
-      end
     end
   end
 end