ruby_llm-infinite_context 0.1.0

data/lib/ruby_llm/infinite_context/pump.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module InfiniteContext
+    # The compactor's scheduler. Builds every node that is not built, has its
+    # sources, and has its whole context summarized (every view line before
+    # the node's end is built). Up to JOBS at once; a failed node is retried
+    # up to `retries` times.
+    module Pump
+      # The nodes to start now, in order: levels from 0 up, and in each
+      # level from the left. `built` and `busy` are sets of [level, i].
+      def self.ready_nodes(total:, first:, built:, busy:, slots:)
+        out = []
+        l = 0
+        while (1 << l) <= total
+          i = 0
+          while ((i + 1) << l) <= total
+            return out if out.size >= slots
+
+            node = [l, i]
+            finish = l.zero? ? i : (i + 1) << l
+            if !built.include?(node) && !busy.include?(node) && finish <= first &&
+                (l.zero? || (built.include?([l - 1, 2 * i]) && built.include?([l - 1, 2 * i + 1])))
+              out << node
+            end
+            i += 1
+          end
+          l += 1
+        end
+        out
+      end
+
+      # Runs a worker thread's body inside the Rails executor when a Rails
+      # app is running, else plain.
+      def self.wrap_thread(&block)
+        executor = (defined?(::Rails) && ::Rails.respond_to?(:application)) ? ::Rails.application&.executor : nil
+        executor ? executor.wrap(&block) : block.call
+      end
+
+      # Runs a blocking wait for worker threads so that they may autoload
+      # constants meanwhile (the Rails load interlock, active when reloading
+      # is on). Plain when ActiveSupport::Dependencies has no interlock.
+      def self.waiting(&block)
+        if defined?(::ActiveSupport::Dependencies) && ::ActiveSupport::Dependencies.respond_to?(:interlock)
+          ::ActiveSupport::Dependencies.interlock.permit_concurrent_loads(&block)
+        else
+          block.call
+        end
+      end
+
+      # Runs passes over the tree until no more nodes can start. Each pass
+      # starts up to `jobs` builder threads and waits for them. `builder`
+      # must answer build(level, i) and return [text, tries].
+      class Runner
+        include TreeMath
+
+        # chat_options: for the default Compactor (ignored when `builder:` is given).
+        def initialize(tree, builder: nil, jobs: JOBS, retries: 3, logger: nil, chat_options: nil)
+          @tree = tree
+          @builder = builder ||
+            Compactor.new(tree, model: tree.compactor_model, chat_options: chat_options)
+          @jobs = jobs
+          @retries = retries
+          @logger = logger
+        end
+
+        # Returns the number of nodes built.
+        #
+        # The workers write on their own connections, so this thread must
+        # read the database fresh: inside a Rails executor (rails runner, a
+        # request, a job) the Active Record query cache is on, and a cached
+        # read of the node sizes never shows a node a worker just stored, so
+        # the same node was built forever (found by bin/install_check). The
+        # cache is bypassed for the run and cleared at the end, so the
+        # caller's next read (the view) is fresh too.
+        def run
+          built = 0
+          attempts = Hash.new(0)
+          ::ActiveRecord::Base.uncached do
+            loop do
+              nodes = startable(Set.new)
+              nodes.reject! { |n| attempts[n] >= @retries }
+              break if nodes.empty?
+
+              run_pass(nodes).each do |node, ok|
+                if ok
+                  built += 1
+                else
+                  attempts[node] += 1
+                end
+              end
+            end
+          end
+          built
+        ensure
+          if ::ActiveRecord::Base.respond_to?(:clear_query_caches_for_current_thread)
+            ::ActiveRecord::Base.clear_query_caches_for_current_thread
+          end
+        end
+
+        private
+
+        def startable(busy)
+          @tree.with_view_lock do
+            sizes = @tree.node_sizes
+            built = sizes.keys.to_set
+            total = @tree.entry_count
+            view = @tree.view(COMPACT_VIEW_BUDGET)
+            first = Fold.new(parts: view.parts) { |l, i| sizes[[l, i]] }.first_unbuilt(total)
+            Pump.ready_nodes(total:, first:, built:, busy:, slots: @jobs)
+          end
+        end
+
+        def run_pass(nodes)
+          done = Queue.new
+          threads = nodes.map do |node|
+            Thread.new do
+              ok = false
+              # In a Rails app each worker thread runs inside the executor (it
+              # returns the thread's database connection when it ends).
+              Pump.wrap_thread do
+                text, tries = @builder.build(*node)
+                @tree.store_node!(*node, text, tries:)
+                ok = true
+              rescue => e
+                log("node #{TreeMath.label(*node)} failed: #{e.class}: #{e.message}")
+              end
+              done << [node, ok]
+            end
+          end
+          # While this thread waits, the workers may autoload constants. With
+          # reloading on (Rails development) that needs the load interlock, so
+          # the wait must permit concurrent loads or the pass deadlocks.
+          Pump.waiting do
+            results = nodes.map { done.pop }
+            threads.each(&:join)
+            results
+          end
+        end
+
+        def log(message)
+          return unless @logger
+
+          (@logger.respond_to?(:call) ? @logger.call(message) : @logger.error("[infinite_context] #{message}"))
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/infinite_context/pump_job.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+if defined?(ActiveJob::Base)
+  module RubyLLM
+    module InfiniteContext
+      # The pump as an ActiveJob: app enqueue it with `set(wait:)` and it
+      # works with any queue backend.
+      class PumpJob < ActiveJob::Base
+        queue_as :default
+
+        def perform(tree_id)
+          tree = Tree.find(tree_id)
+          Pump::Runner.new(tree, logger: Rails.logger).run
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/infinite_context/schema.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module InfiniteContext
+    # The tables, created by the install generator's migration. `json`
+    # columns (not jsonb) and row locks keep SQLite and Postgres both working.
+    module Schema
+      TABLES = %w[ric_trees ric_entries ric_nodes ric_views ric_zoom_events].freeze
+
+      def self.migration(migrator)
+        migrator.create_table :ric_trees do |t|
+          t.string :key, null: false
+          t.string :compactor_model
+          t.json :options, null: false, default: {}
+          t.timestamps
+        end
+        migrator.add_index :ric_trees, :key, unique: true
+
+        # The append-only log. Never updated or deleted.
+        migrator.create_table :ric_entries do |t|
+          t.references :tree, null: false, index: false
+          t.integer :idx, null: false
+          t.string :kind, null: false
+          t.text :text, null: false
+          t.integer :bytes, null: false
+          t.datetime :created_at, null: false
+        end
+        migrator.add_index :ric_entries, [:tree_id, :idx], unique: true
+
+        # The summary tree. Node (level, i) covers entries
+        # [i * 2^level, (i + 1) * 2^level). Built once, never recomputed.
+        # dig_count and last_dug_at count how often the agent opened it.
+        migrator.create_table :ric_nodes do |t|
+          t.references :tree, null: false, index: false
+          t.integer :level, null: false
+          t.integer :i, null: false
+          t.text :text, null: false
+          t.integer :bytes, null: false
+          t.integer :tries, null: false, default: 0
+          t.integer :dig_count, null: false, default: 0
+          t.datetime :last_dug_at
+          t.datetime :created_at, null: false
+        end
+        migrator.add_index :ric_nodes, [:tree_id, :level, :i], unique: true
+
+        # One saved view per budget.
+        migrator.create_table :ric_views do |t|
+          t.references :tree, null: false, index: false
+          t.integer :budget, null: false
+          t.json :parts, null: false, default: []
+          t.integer :message_count, null: false, default: 0
+          t.timestamps
+        end
+        migrator.add_index :ric_views, [:tree_id, :budget], unique: true
+
+        # One row per zoom: the line opened, the leaf the dig reached (when
+        # it went all the way down), the steps it cost, and the turn.
+        migrator.create_table :ric_zoom_events do |t|
+          t.references :tree, null: false, index: false
+          t.string :session, null: false
+          t.integer :level, null: false
+          t.integer :i, null: false
+          t.integer :leaf_idx
+          t.integer :steps, null: false, default: 0
+          t.datetime :created_at, null: false
+        end
+        migrator.add_index :ric_zoom_events, [:tree_id, :session]
+        migrator.add_index :ric_zoom_events, [:tree_id, :level, :i]
+      end
+
+      # For tests and for apps that manage their own schema (no migration
+      # file needed).
+      def self.create!(connection = ::ActiveRecord::Base.connection)
+        shim = Object.new
+        shim.define_singleton_method(:create_table) { |*a, **o, &b| connection.create_table(*a, **o, &b) }
+        shim.define_singleton_method(:add_index) { |*a, **o| connection.add_index(*a, **o) }
+        migration(shim)
+      end
+    end
+  end
+end

data/lib/ruby_llm/infinite_context/tokens.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module InfiniteContext
+    # Budgets are UTF-8 bytes everywhere inside the gem. Callers think in
+    # tokens, so two doors:
+    #
+    # - `20_000.tokens` via the refinement (`using RubyLLM::InfiniteContext::Tokens`),
+    # - a `tokens:` keyword on attach, with_infinite_context and has_infinite_context.
+    #
+    # The estimator is a rough 4 bytes per token, right for English prose;
+    # log text with code or non-ASCII runs heavier. Callers who need exact
+    # token counts should pass `budget:` in bytes instead.
+    BYTES_PER_TOKEN = 4
+
+    def self.token_bytes(tokens)
+      tokens * BYTES_PER_TOKEN
+    end
+
+    module Tokens
+      refine Integer do
+        def tokens = RubyLLM::InfiniteContext.token_bytes(self)
+      end
+    end
+
+    # Resolves the budget from (bytes, tokens) arguments; bytes win when
+    # both are given.
+    def self.resolve_budget(budget: nil, tokens: nil)
+      if budget
+        budget
+      elsif tokens
+        token_bytes(tokens)
+      else
+        DEFAULT_VIEW_BUDGET
+      end
+    end
+  end
+end

data/lib/ruby_llm/infinite_context/tools.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module InfiniteContext
+    # The agent's tools, as RubyLLM::Tool classes. Instances carry the tree
+    # and the turn, so one Turn guards and caps the zooms of one reply.
+    require "ruby_llm/tool"
+
+    module Tools
+      # zoom(id, n): open the line covering messages id..id+n-1 into the two
+      # lines it was built from; zoom(id, 1) shows message id itself.
+      class ZoomTool < ::RubyLLM::Tool
+        # RubyLLM derives a name from the namespaced class
+        # (ruby_llm--infinite_context--tools--zoom), but the prompts and the
+        # view guide say "zoom". The model must see the name we teach it.
+        def self.tool_name = "zoom"
+
+        description "Open a line of the conversation view. zoom(id, n) returns the two " \
+                    "half-size lines the line covering messages id..id+n-1 was built from. " \
+                    "zoom(id, 1) shows message id in full. Zoom before you guess: when a " \
+                    "summary only mentions something you need, open the line that mentions it."
+        parameter :id, type: "integer", description: "Start message index of the line"
+        parameter :n, type: "integer", description: "Number of messages the line covers (a power of 2); 1 for a single message"
+
+        def initialize(handler)
+          @handler = handler
+        end
+
+        def execute(id:, n:) = @handler.zoom(id, n)
+      end
+
+      # read(id, page): the next page of a long message.
+      class ReadTool < ::RubyLLM::Tool
+        def self.tool_name = "read"
+
+        description "Read a further page of a long message. zoom(id, 1) shows its first page."
+        parameter :id, type: "integer", description: "Message index"
+        parameter :page, type: "integer", description: "Page number, starting at 2"
+
+        def initialize(handler)
+          @handler = handler
+        end
+
+        def execute(id:, page: 2) = @handler.read(id, page)
+      end
+
+      # date(id): when a message was written.
+      class DateTool < ::RubyLLM::Tool
+        def self.tool_name = "date"
+
+        description "Date and time a message was written."
+        parameter :id, type: "integer", description: "Message index"
+
+        def initialize(handler)
+          @handler = handler
+        end
+
+        def execute(id:) = @handler.date(id)
+      end
+
+      # The tool set for one turn: fresh instances per turn, because the
+      # guard and the cap are per turn.
+      def self.for_turn(zoom_handler, turn)
+        [ZoomTool.new(Zoom.new(zoom_handler, turn)),
+          ReadTool.new(Zoom.new(zoom_handler, turn)),
+          DateTool.new(Zoom.new(zoom_handler, turn))]
+      end
+    end
+  end
+end

data/lib/ruby_llm/infinite_context/tree_math.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module InfiniteContext
+    # Shared node math and byte helpers.
+    #
+    # Node (level, i) covers entries [first, last) of the append-only log,
+    # where first = i << level and last = (i + 1) << level. Its label,
+    # "id+n", names the stretch in the view and in the zoom tools.
+    module TreeMath
+      module_function
+
+      def first(level, i) = i << level
+
+      def last(level, i) = (i + 1) << level
+
+      def label(level, i) = "#{first(level, i)}+#{1 << level}"
+
+      def bytes(text) = text.to_s.bytesize
+
+      # First `limit` bytes of text, never splitting a UTF-8 character.
+      def cut(text, limit)
+        text.byteslice(0, limit).scrub("").delete_suffix("\uFFFD")
+      end
+
+      # One line: whitespace runs around a newline become one space.
+      def flat(text) = text.to_s.gsub(/\s*\n\s*/, " ")
+
+      # Head and tail of an over-long text, with a note of what was cut.
+      def cap(text, limit = CAP)
+        return text if text.length <= limit
+
+        half = limit / 2
+        cut_chars = text.length - 2 * half
+        "#{text[0, half]}\n[... #{cut_chars} characters cut ...]\n#{text[-half..]}"
+      end
+    end
+  end
+end

data/lib/ruby_llm/infinite_context/version.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module InfiniteContext
+    VERSION = "0.1.0"
+  end
+end

data/lib/ruby_llm/infinite_context/zoom.rb

@@ -0,0 +1,135 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module InfiniteContext
+    # Per-turn state for the zoom tools: the lines opened so far, the zoom
+    # count, and the turn's id (the `session` of a zoom event).
+    class Turn
+      attr_reader :id, :opened
+
+      def initialize(id = nil, cap: ZOOM_CAP_PER_TURN)
+        @id = id || SecureRandom.uuid
+        @cap = cap
+        @opened = {}
+        @zooms = 0
+      end
+
+      # True when the agent opened this label earlier in the turn.
+      def repeat?(label) = @opened.key?(label)
+
+      def open!(label, text)
+        @opened[label] = text
+        @zooms += 1
+      end
+
+      def cap_reached? = @zooms >= @cap
+      def remaining = @cap - @zooms
+
+      attr_reader :zooms
+    end
+
+    # The agent's memory tools: zoom, read (paging through a long message)
+    # and date. `#zoom` refuses repeats ("you already opened this line") and
+    # enforces the per-turn cap, and records counters and zoom events.
+    class Zoom
+      include TreeMath
+
+      def initialize(tree, turn)
+        @tree = tree
+        @turn = turn
+      end
+
+      def zoom(id, n)
+        id = Integer(id)
+        n = Integer(n)
+        total = @tree.entry_count
+        return "No line #{id}+#{n}." unless valid?(id, n, total)
+
+        return leaf(id) if n == 1
+
+        label = TreeMath.label(level_of(n), left_of(id, n))
+        return already_opened(label) if @turn.repeat?(label)
+        return cap_message if @turn.cap_reached?
+
+        @turn.open!(label, true)
+        level = level_of(n)
+        left = left_of(id, n)
+        text = [left, left + 1].map do |i|
+          child = @tree.node_text(level, i)
+          child ? "#{TreeMath.label(level, i)}|#{TreeMath.flat(child)}" : "#{TreeMath.label(level, i)}|#{PLACEHOLDER}"
+        end.join("\n")
+        @tree.record_zoom!(level:, i: left, leaf_idx: nil, steps: @turn.zooms, session: @turn.id)
+        text
+      rescue ArgumentError, TypeError
+        "No line #{id}+#{n}."
+      end
+
+      # A page of one message, for messages longer than one page.
+      def read(id, page = 1)
+        id = Integer(id)
+        page = Integer(page)
+        entry = @tree.entry(id)
+        return "No message #{id}." unless entry && page.positive?
+
+        text = entry.text.to_s
+        pages = (TreeMath.bytes(text).to_f / PAGE).ceil
+        return "Message #{id} has 1 page; it is shown in full by zoom(#{id}, 1)." if pages <= 1 || page > pages
+
+        slice = text.byteslice((page - 1) * PAGE, PAGE).to_s
+        suffix = (page < pages) ? "\n[page #{page} of #{pages}; read(#{id}, #{page + 1}) for more]" : "\n[end of message #{id}]"
+        "#{slice}#{suffix}"
+      rescue ArgumentError, TypeError
+        "No message #{id}."
+      end
+
+      def date(id)
+        entry = @tree.entry(Integer(id))
+        return "No message #{id}." unless entry
+
+        entry.created_at.strftime("%a %Y-%m-%d %H:%M %Z")
+      rescue ArgumentError, TypeError
+        "No message #{id}."
+      end
+
+      private
+
+      # One message in full (paged when long). Reaching the leaf is what the
+      # counters count: the leaf reached and the steps the dig cost.
+      def leaf(id)
+        label = TreeMath.label(0, id)
+        return already_opened(label) if @turn.repeat?(label)
+        return cap_message if @turn.cap_reached?
+
+        entry = @tree.entry(id)
+        return "No line #{id}+1." unless entry
+
+        @turn.open!(label, true)
+        @tree.record_zoom!(level: 0, i: id, leaf_idx: id, steps: @turn.zooms, session: @turn.id)
+        text = entry.text.to_s
+        if TreeMath.bytes(text) > PAGE
+          pages = (TreeMath.bytes(text).to_f / PAGE).ceil
+          "#{text.byteslice(0, PAGE)}\n[page 1 of #{pages}; read(#{id}, 2) for more]"
+        else
+          text
+        end
+      end
+
+      def already_opened(label)
+        "You already opened line #{label} this turn; its text is above. Use what you have."
+      end
+
+      def cap_message
+        "Zoom cap reached for this turn (#{@turn.opened.size} lines opened). " \
+          "Work with the lines you have and say what you still need in your reply."
+      end
+
+      def valid?(id, n, total)
+        n.positive? && (n & (n - 1)).zero? && (id % n).zero? && id >= 0 && id + n <= total
+      end
+
+      def level_of(n) = n.bit_length - 2
+
+      def left_of(id, n) = 2 * id / n
+    end
+  end
+end

data/lib/ruby_llm/infinite_context.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+# ruby_llm-infinite_context: endless memory for RubyLLM 2.x chats and
+# agents. Every message is logged, a cheap model compresses the log into a
+# binary tree of one-line summaries, and the chat sees a saved view of the
+# tree that fits a byte budget the caller picks. The agent gets tools to
+# zoom into any line, down to the original message.
+#
+# Based on Victor Taelin's OptChat/UniiChat design
+# (https://gist.github.com/VictorTaelin/91837951a5ce5b38f341ec1ba1df6449)
+# and on pi-optchat (https://github.com/jonaslsaa/pi-optchat). See README.md.
+
+require_relative "infinite_context/version"
+
+require "ruby_llm" # the full library: RubyLLM.chat and friends
+
+module RubyLLM
+  module InfiniteContext
+    class Error < StandardError; end
+
+    NODE = 512              # target size of one summary line
+    DEFAULT_VIEW_BUDGET = 128_000 # default consumer view budget, bytes
+    COMPACT_VIEW_BUDGET = 24_000  # the compactor's own view, 16-32 KB
+    JOBS = 8                # compactor calls at once
+    MARGIN = 64             # bytes a non-Claude compactor is asked to stay under NODE
+    TRIES = 5               # attempts per node to get under NODE
+    CAP = 30_000            # max characters of one stubbed tool result (head and tail)
+    PAGE = 4_000            # bytes per page when reading a long message
+    PLACEHOLDER = "(not summarized yet: zoom it)"
+    ZOOM_CAP_PER_TURN = 12  # max lines the agent may open in one turn
+    KINDS = %w[user assistant tool echo note].freeze
+  end
+end
+
+require_relative "infinite_context/tree_math"
+require_relative "infinite_context/schema"
+require_relative "infinite_context/tokens"
+require_relative "infinite_context/fold"
+require_relative "infinite_context/prompts"
+require_relative "infinite_context/chat_options"
+require_relative "infinite_context/compactor"
+require_relative "infinite_context/pump"
+require_relative "infinite_context/cache_mark"
+require_relative "infinite_context/zoom"
+require_relative "infinite_context/tools"
+require_relative "infinite_context/batch_import"
+require_relative "infinite_context/manager"
+
+begin
+  require "rails"
+  require_relative "infinite_context/models"
+  require_relative "infinite_context/pump_job"
+  require_relative "infinite_context/model"
+  require_relative "infinite_context/engine"
+rescue LoadError
+  # Not a Rails app: the Active Record models load lazily on first use.
+end