ruby_llm 2.1.1 → 2.2.0

data/lib/ruby_llm/active_record/claim_heartbeat.rb

@@ -0,0 +1,149 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module ActiveRecord
+    # Calls its block every +interval+ seconds from a thread of its own
+    # until stopped. The thread connects as the starting thread did, to the
+    # same roles and shards, and each beat checks out a connection from the
+    # pool of each of +models+ and returns it, so the thread holds none
+    # between beats.
+    #
+    # A beat that fails is skipped, and the next one tries again. A lease
+    # spans at least four intervals, so a claim outlasts a skipped beat,
+    # but after two in a row it may lapse before the next one lands, so
+    # the second logs a warning.
+    #
+    # Once its process is about to exit, ::hurry_all makes every live
+    # heartbeat beat at once and then every HURRIED_INTERVAL seconds, so a
+    # claim a run wrote just before learning of the exit is cut back too.
+    class ClaimHeartbeat # :nodoc:
+      HURRIED_INTERVAL = 1
+
+      LIVE_LOCK = Mutex.new
+      private_constant :LIVE_LOCK
+
+      @live = Set.new
+
+      class << self
+        def hurry_all
+          LIVE_LOCK.synchronize { @live.to_a }.each(&:hurry)
+        end
+
+        def add(heartbeat)
+          LIVE_LOCK.synchronize { @live << heartbeat }
+        end
+
+        def delete(heartbeat)
+          LIVE_LOCK.synchronize { @live.delete(heartbeat) }
+        end
+      end
+
+      def initialize(models, interval, &beat)
+        @models = models
+        @interval = interval
+        @beat = beat
+        @failed_beats = 0
+        @lock = Mutex.new
+        @wake = ConditionVariable.new
+        @stopped = false
+        @hurried = false
+        @thread = start(::ActiveRecord::Base.connected_to_stack.to_a)
+        self.class.add(self)
+      end
+
+      def stop
+        self.class.delete(self)
+        @lock.synchronize do
+          @stopped = true
+          @wake.signal
+        end
+        @thread.join
+      end
+
+      def hurry
+        @lock.synchronize do
+          @interval = [@interval, HURRIED_INTERVAL].min
+          @hurried = true
+          @wake.signal
+        end
+      end
+
+      private
+
+      def start(connected_to)
+        thread = Thread.new do
+          ::ActiveRecord::Base.connected_to_stack.concat(connected_to)
+          beat until stopped_after_interval?
+        end
+        thread.name = 'ruby_llm claim heartbeat'
+        thread.report_on_exception = false
+        thread
+      end
+
+      def stopped_after_interval?
+        @lock.synchronize do
+          @wake.wait(@lock, @interval) unless @stopped || @hurried
+          @hurried = false
+          @stopped
+        end
+      end
+
+      def beat
+        deadline = monotonic_now + (@interval / 4.0)
+        connected(@models.map(&:connection_pool).uniq, deadline) { @beat.call }
+        @failed_beats = 0
+      rescue StandardError => e
+        beat_failed(e)
+      end
+
+      def beat_failed(error)
+        @failed_beats += 1
+        if @failed_beats == 1 && busy?(error)
+          RubyLLM.logger.debug("RubyLLM skipped renewing a chat's claims, the database was busy")
+        else
+          RubyLLM.logger.warn("RubyLLM could not renew a chat's claims (#{@failed_beats} in a row): " \
+                              "#{error.class}: #{error.message}")
+        end
+      end
+
+      def busy?(error)
+        [error, error.cause].any?(::SQLite3::BusyException) if defined?(::SQLite3::BusyException)
+      end
+
+      def connected(pools, deadline, &)
+        return yield if pools.empty?
+
+        pools.first.with_connection do |connection|
+          waiting_until(connection, deadline) { connected(pools.drop(1), deadline, &) }
+        end
+      end
+
+      # Before Rails 8, SQLite waits for a lock through busy_timeout, which
+      # holds the GVL and stops every thread in the process, including the
+      # run whose transaction holds the lock. The beat waits with the GVL
+      # released instead, until its deadline, and then disconnects, so the
+      # next query on that connection reconnects with the application's own
+      # busy settings. An in-memory database has no other connection to wait
+      # for, and disconnecting would drop it.
+      def waiting_until(connection, deadline)
+        return yield unless gvl_held_while_busy?(connection)
+
+        begin
+          connection.raw_connection.busy_handler { monotonic_now < deadline && sleep(0.001) }
+          yield
+        ensure
+          connection.disconnect!
+        end
+      end
+
+      def gvl_held_while_busy?(connection)
+        connection.adapter_name == 'SQLite' && ::ActiveRecord::VERSION::MAJOR < 8 &&
+          connection.supports_concurrent_connections?
+      end
+
+      def monotonic_now
+        Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      end
+    end
+  end
+end

data/lib/ruby_llm/active_record/file_store.rb

@@ -0,0 +1,269 @@
+# frozen_string_literal: true
+
+require 'stringio'
+require 'tempfile'
+
+module RubyLLM
+  module ActiveRecord
+    # Keeps a persisted chat's files in +ruby_llm_chat_files+, with their
+    # content in Active Storage, so they live on whatever service Active
+    # Storage uses: a disk now, a bucket later. A persisted chat uses one
+    # when ChatFiles#with_file_store is given no store.
+    #
+    #   chat.with_file_store
+    #   chat.file_store.list
+    #
+    # The store writes files for the chat's owner, the record
+    # <tt>acts_as_chat owner:</tt> names or one set with
+    # ChatMethods#with_owner, and sees only their files and the files
+    # without an owner, so in a chat several people share, one person's
+    # tool results stay theirs:
+    #
+    #   chat.with_owner(Current.user).with_file_store
+    #
+    # What the model last read of a file is kept in the row's
+    # +read_digest+ column, so it holds across processes, and #with_lock
+    # locks the chat's row. Converted text is kept in the row's +text+
+    # column. Encrypt it, like
+    # anything else the model reads, with
+    # <tt>RubyLLM::ActiveRecord::ChatFile.encrypts :text</tt>.
+    class FileStore < RubyLLM::FileStore
+      # The chat record the files belong to.
+      attr_reader :chat
+
+      # Creates a store for the files of the +chat+ record, writing them for
+      # +owner+ when given, and otherwise for the chat's owner.
+      #
+      # Raises ArgumentError when +chat+ is not an +acts_as_chat+ record,
+      # and ConfigurationError when Active Storage is not loaded.
+      def initialize(chat, owner: nil)
+        unless chat.is_a?(ChatFiles)
+          raise ArgumentError, "#{self.class} needs an acts_as_chat record, not #{chat.class}"
+        end
+        unless ChatFile.respond_to?(:reflect_on_attachment) && ChatFile.reflect_on_attachment(:file)
+          raise ConfigurationError, "#{self.class} needs Active Storage. " \
+                                    'Load it with require "active_storage/engine" in config/application.rb, ' \
+                                    'then run bin/rails active_storage:install and bin/rails db:migrate.'
+        end
+
+        super()
+        @chat = chat
+        @owner = owner
+      end
+
+      # Returns the owner the store writes files for: the one it was
+      # created with, or else the chat's owner when that is a record.
+      # Returns +nil+ when there is neither.
+      def owner
+        return @owner unless @owner.nil?
+
+        owner = @chat.to_llm.owner
+        owner if owner.is_a?(::ActiveRecord::Base)
+      end
+
+      # Returns the files inside the folder +prefix+, or every file.
+      def list(prefix = nil)
+        scope = prefix.present? ? inside(path(prefix)) : visible
+        preferred(scope.includes(:owner).order(:path, :id).to_a).map(&:to_stored_file)
+      end
+
+      # Returns the StoredFile at +path+, or +nil+.
+      def [](path)
+        row(path)&.to_stored_file
+      end
+
+      # Returns the content of the file at +path+.
+      def read(path)
+        record = row(path) || raise(FileNotFoundError, "No file at #{path(path)}")
+        record.file.download.b
+      end
+
+      # Writes +content+ to +path+ and returns its StoredFile.
+      def write(path, content, content_type: nil)
+        path = path(path)
+        rewindable(content) do |io|
+          type = content_type || RubyLLM::FileStore.content_type_for(path, io)
+          digest = RubyLLM::FileStore.digest(io)
+          io.rewind
+          blob = ::ActiveStorage::Blob.create_and_upload!(io:, filename: File.basename(path),
+                                                          content_type: type, identify: false)
+          save(path, blob.byte_size, type, digest) { |record| record.file.attach(blob) }
+        end
+      end
+
+      # Deletes the file at +path+ and returns its StoredFile, or +nil+. A
+      # store with an owner deletes only that owner's files.
+      def delete(path)
+        path = path(path)
+        with_lock do
+          record = owned.find_by(path:)
+          refuse_folder(path) if record.nil? && rows_in_folder(path).any? { |row| row.path != path }
+          record&.destroy
+          record&.to_stored_file
+        end
+      end
+
+      # Deletes everything inside the folder +path+ and returns the deleted
+      # StoredFile objects. A store with an owner deletes only that
+      # owner's files.
+      def delete_folder(path)
+        path = path(path)
+        with_lock do
+          refuse_file(path) if row(path)
+
+          rows_in_folder(path, owned).each(&:destroy).map(&:to_stored_file)
+        end
+      end
+
+      # Returns the digest of what the model last read at +path+, from the
+      # file's row, or +nil+.
+      def read_digest(path) # :nodoc:
+        return super unless read_digests?
+
+        row(path)&.read_digest
+      end
+
+      # Keeps +digest+ as what the model read at +path+, in the file's row.
+      def record_read(path, digest) # :nodoc:
+        return super unless read_digests?
+
+        row(path)&.update_columns(read_digest: digest)
+      end
+
+      # Runs the block in a transaction that locks the chat's row, past
+      # the query cache, so other processes wait for it.
+      def with_lock
+        super do
+          ChatFile.uncached do
+            ChatFile.transaction do
+              lock_chat
+              yield
+            end
+          end
+        end
+      end
+
+      # Returns the text kept for +file+, or +nil+.
+      def cached_text(file) # :nodoc:
+        record = row(file.path)
+        record.text if record&.digest == file.digest && record.has_attribute?(:text)
+      end
+
+      # Keeps +text+ for +file+ in its row.
+      def cache_text(file, text) # :nodoc:
+        record = row(file.path)
+        record.update_columns(text:) if record&.digest == file.digest && record.has_attribute?(:text)
+      end
+
+      # Keeps an attachment from Active Storage by its blob, without
+      # copying it: the file and the message share the blob.
+      def store_attachment(attachment) # :nodoc:
+        blob = blob_of(attachment)
+        return super unless blob
+
+        existing = visible.joins(:file_blob).find_by(active_storage_blobs: { id: blob.id })
+        existing ? existing.to_stored_file : super
+      end
+
+      private
+
+      # Active Storage reads an IO from its start, so one that cannot
+      # rewind, or stands past its start, is copied to a tempfile first.
+      def rewindable(content)
+        return yield StringIO.new(content.to_s.b) unless content.respond_to?(:read)
+        return yield content if at_start?(content)
+
+        Tempfile.create('ruby_llm-file-', binmode: true) do |copy|
+          IO.copy_stream(content, copy)
+          copy.rewind
+          yield copy
+        end
+      end
+
+      def at_start?(io)
+        %i[rewind pos].all? { |method| io.respond_to?(method) } && io.pos.zero?
+      rescue Errno::ESPIPE
+        false
+      end
+
+      def read_digests?
+        ChatFile.column_names.include?('read_digest')
+      end
+
+      def rows_in_folder(folder, scope = visible)
+        inside(folder, scope).includes(:owner).order(:path, :id)
+                             .select { |row| RubyLLM::FileStore.inside?(row.path, folder) }
+      end
+
+      def write_attachment(path, attachment, content)
+        blob = blob_of(attachment)
+        return super unless blob
+
+        save(path, blob.byte_size, attachment.mime_type, RubyLLM::FileStore.digest(content)) do |record|
+          record.file.attach(blob)
+        end
+      end
+
+      def save(path, byte_size, content_type, digest)
+        ChatFile.transaction do
+          lock_chat
+          record = owned.lock.find_or_initialize_by(path:)
+          record.assign_attributes(byte_size:, content_type:, digest:)
+          record.text = nil if record.has_attribute?(:text)
+          yield record
+          record.save!
+          record.to_stored_file
+        end
+      end
+
+      # SQLite has no row locks, and a deferred transaction that reads
+      # first fails when it later writes beside another writer, so it takes
+      # the write lock at once with a no-op update.
+      def lock_chat
+        key = @chat.class.primary_key
+        scope = @chat.class.unscoped.where(key => @chat.id)
+        return scope.lock.pick(key) unless @chat.class.connection.adapter_name.match?(/sqlite/i)
+
+        column = @chat.class.connection.quote_column_name(key)
+        scope.update_all("#{column} = #{column}")
+      end
+
+      def row(path)
+        preferred(visible.where(path: path(path)).order(:id).to_a).first
+      end
+
+      def preferred(records)
+        records.group_by(&:path).values.map { |same| same.find(&:owner_id) || same.first }
+      end
+
+      def rows
+        ChatFile.where(chat: @chat)
+      end
+
+      def visible
+        shared = rows.where(owner_type: nil, owner_id: nil)
+        owner ? shared.or(rows.where(owner:)) : shared
+      end
+
+      def owned
+        owner ? rows.where(owner:) : rows.where(owner_type: nil, owner_id: nil)
+      end
+
+      def inside(folder, scope = visible)
+        inside = "#{ChatFile.sanitize_sql_like(folder, '!')}/%"
+        scope.where("path = :folder OR path LIKE :inside ESCAPE '!'", folder:, inside:)
+      end
+
+      def path(name)
+        RubyLLM::FileStore.path(name)
+      end
+
+      def blob_of(attachment)
+        source = attachment.source
+        return source if defined?(::ActiveStorage::Blob) && source.is_a?(::ActiveStorage::Blob)
+
+        source.blob if source.respond_to?(:blob) && source.blob.is_a?(::ActiveStorage::Blob)
+      end
+    end
+  end
+end

data/lib/ruby_llm/active_record/mcp_credential.rb

@@ -2,9 +2,23 @@
 
 module RubyLLM
   module ActiveRecord
-    # RubyLLM's persistence for MCP OAuth credentials, encrypted with Active
-    # Record encryption.
-    class MCPCredential < Record # :nodoc:
+    # The OAuth state RubyLLM::MCP keeps for a server, in RubyLLM's
+    # +ruby_llm_mcp_credentials+ table. +data+ is encrypted with Active
+    # Record encryption, so the application needs encryption keys.
+    #
+    #   RubyLLM::ActiveRecord::MCPCredential.where(owner: current_user).pluck(:updated_at)
+    #
+    # Each row has an optional polymorphic +owner+, the record the
+    # credentials were authorized for. Its columns:
+    #
+    # [+key+] The key RubyLLM stores the state under, unique across the table.
+    # [+data+] Tokens, client registrations, and signing keys, as a Hash.
+    #          Treat it as a secret.
+    #
+    # RubyLLM writes these rows. Read and query them, and add scopes or
+    # columns of your own, but do not create, update, or delete them: call
+    # +deauthorize+ on the server to forget a user's credentials.
+    class MCPCredential < Record
       self.table_name = 'ruby_llm_mcp_credentials'
 
       belongs_to :owner, polymorphic: true, optional: true
@@ -15,21 +29,22 @@ module RubyLLM
       validates :key, presence: true
 
       class << self
-        def read(key)
+        def read(key) # :nodoc:
           find_by(key:)&.data
         end
 
-        def write(key, data, owner: nil)
-          record = find_or_initialize_by(key:)
-          record.update!(data:, owner: owner.is_a?(::ActiveRecord::Base) ? owner : nil)
+        def write(key, data, owner: nil) # :nodoc:
+          attributes = { data:, owner: owner.is_a?(::ActiveRecord::Base) ? owner : nil }
+          record = find_by(key:) || create_or_find_by!(key:) { |created| created.assign_attributes(attributes) }
+          record.update!(attributes)
         end
 
-        def delete(key)
+        def delete(key) # :nodoc:
           where(key:).delete_all
         end
 
         # Updating the row locks it on every adapter; SQLite ignores FOR UPDATE.
-        def synchronize(key)
+        def synchronize(key) # :nodoc:
           transaction do
             where(key:).touch_all
             yield

data/lib/ruby_llm/active_record/message_broadcasting.rb

@@ -0,0 +1,85 @@
+# frozen_string_literal: true
+
+require 'active_support/concern'
+
+module RubyLLM
+  module ActiveRecord
+    # Shows a chat's messages live with Turbo when
+    # Configuration#broadcast_messages, or the model's +broadcasts:+
+    # option, is on. An empty assistant reply is
+    # appended when it is created, so its chunks have a place to land, and
+    # each chunk is appended as a numbered span that the chat UI's Stimulus
+    # controller puts back in order. A saved message replaces itself, a
+    # destroyed one is removed, and every change refreshes the chat page.
+    # A refresh carries no request id, so it reaches the page whose request
+    # made the change too, and Turbo never skips a burst of refreshes as
+    # that page's own. Broadcasts go to the chat record's stream:
+    #
+    #   <%= turbo_stream_from @chat %>
+    #
+    module MessageBroadcasting
+      extend ActiveSupport::Concern
+
+      # Returns whether Turbo broadcasting can run: turbo-rails and Action
+      # Cable are loaded.
+      def self.available?
+        defined?(::Turbo::Broadcastable) && defined?(::ActionCable) ? true : false
+      end
+
+      included do
+        after_create_commit -> { broadcast_append_to chat_association }, if: :streaming_placeholder?
+        after_update_commit -> { broadcast_replace_to chat_association }, if: :broadcasting?
+        after_destroy_commit -> { broadcast_remove_to chat_association }, if: :broadcasting?
+        after_commit -> { broadcast_refresh_to chat_association, request_id: nil }, if: :broadcasting?
+      end
+
+      class_methods do
+        # Returns whether this model's messages broadcast: its
+        # +broadcasts:+ option when given, Configuration#broadcast_messages
+        # otherwise.
+        def broadcasting?
+          ruby_llm_broadcasts.nil? ? RubyLLM.config.broadcast_messages == true : ruby_llm_broadcasts
+        end
+      end
+
+      # Returns whether this message broadcasts. See ::broadcasting?.
+      def broadcasting?
+        self.class.broadcasting?
+      end
+
+      # Appends the content of +chunk+, a RubyLLM::Chunk streamed into this
+      # message, to the message's <tt><dom_id>_stream</tt> element as
+      # <tt><span data-seq="index"></tt>, escaped. Each span names the
+      # index of the span before it in <tt>data-after</tt>, so the page can
+      # order them. A chunk can reach the page before the message does, and
+      # Turbo drops an append with no target, so a page of your own must hold
+      # it, as the chat UI's Stimulus controller does. Returns +nil+ for a
+      # chunk without content. Raises
+      # RubyLLM::ConfigurationError when the message does not broadcast.
+      #
+      #   chat.ask("Tell me a story") { |chunk| chunk.message.broadcast_chunk(chunk) }
+      #
+      def broadcast_chunk(chunk)
+        raise RubyLLM::ConfigurationError, ChunkMethods::BROADCASTING_OFF unless broadcasting?
+        return if chunk.content.to_s.empty?
+
+        after = @last_broadcast_chunk_index
+        @last_broadcast_chunk_index = chunk.index
+        broadcast_append_to chat_association,
+                            target: "#{ActionView::RecordIdentifier.dom_id(self)}_stream",
+                            content: chunk_span(chunk, after)
+      end
+
+      private
+
+      def streaming_placeholder?
+        broadcasting? && role.to_s == 'assistant' && content.blank?
+      end
+
+      def chunk_span(chunk, after)
+        after_attribute = %( data-after="#{after}") if after
+        %(<span data-seq="#{chunk.index}"#{after_attribute}>#{ERB::Util.html_escape(chunk.content)}</span>)
+      end
+    end
+  end
+end