ruby_llm 2.0.0 → 2.1.0

data/lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb

@@ -1,341 +0,0 @@
-# frozen_string_literal: true
-
-require 'json'
-require_relative '../legacy_content_sql'
-require_relative 'verification'
-
-module RubyLLM
-  module Generators
-    class OnlineCopyMigration < UpgradeMigration
-      class Data < UpgradeMigration # :nodoc: all
-        BATCH_SIZE = 10_000
-        PROGRESS = :ruby_llm_v2_backfills
-        TASKS = %w[message_content tool_results usages].freeze
-
-        def initialize(connection:, settings:)
-          super(connection:)
-          @configuration = settings
-        end
-
-        def prepare
-          messages = configuration.fetch('message_table')
-          unless @connection.column_exists?(messages, :ruby_llm_content)
-            @connection.add_column(messages, :ruby_llm_content, :text)
-          end
-          %i[ruby_llm_tool_calls ruby_llm_usages].each do |table|
-            @connection.add_column(table, :legacy_key, :string) unless @connection.column_exists?(table, :legacy_key)
-            unless @connection.index_exists?(table, :legacy_key, unique: true)
-              @connection.add_index(table, :legacy_key, unique: true)
-            end
-          end
-          return if @connection.column_exists?(:ruby_llm_tool_calls, :legacy_chat_id)
-
-          @connection.add_column(:ruby_llm_tool_calls, :legacy_chat_id, :string)
-          @connection.add_index(:ruby_llm_tool_calls, :legacy_chat_id)
-        end
-
-        def backfill
-          prepare_progress
-          sync_models
-          sync_chat_models
-          TASKS.each do |task|
-            progress = records(PROGRESS).find_or_create_by!(task:)
-            next if progress.completed
-
-            unless progress.upper_id
-              upper = messages.order(messages.primary_key => :desc).pick(messages.primary_key)
-              unless upper
-                progress.update!(completed: true)
-                next
-              end
-              progress.update!(upper_id: upper&.to_s)
-            end
-            relation = legacy_messages
-            relation = relation.where("#{q(messages.primary_key)} <= ?", progress.upper_id) if progress.upper_id
-            relation = relation.where("#{q(messages.primary_key)} > ?", progress.last_id) if progress.last_id
-            each_batch(relation) do |batch|
-              messages.transaction do
-                copy_batch(task, batch)
-                progress.update!(last_id: batch.last.id)
-              end
-            end
-            progress.update!(completed: true)
-          end
-        end
-
-        def discard_incomplete_tool_calls
-          source = records(configuration.fetch('tool_call_table'))
-          result_key = configuration.fetch('tool_call_foreign_key')
-          results = messages.where.not(result_key => nil).select(result_key)
-          owners = legacy_messages.select(messages.primary_key)
-          source.where(configuration.fetch('message_foreign_key') => owners)
-                .where.not(source.primary_key => results).delete_all
-        end
-
-        def catch_up(passes: nil)
-          sync_models
-          pass = 0
-          loop do
-            upper = changes.maximum(:id)
-            break unless upper
-
-            changes.where(id: ..upper).in_batches(of: 100) do |events|
-              captured = events.to_a
-              chat_ids = affected_chats(captured)
-              sync_chat_models(chat_ids)
-              reconcile_chats(chat_ids)
-              acknowledge(captured)
-            end
-            pass += 1
-            break if passes && pass >= passes
-          end
-        end
-
-        def completed?
-          @connection.table_exists?(PROGRESS) &&
-            (TASKS - records(PROGRESS).where(completed: true).pluck(:task)).empty?
-        end
-
-        def finished?
-          @connection.table_exists?(PROGRESS) && records(PROGRESS).where(task: 'finished', completed: true).exists?
-        end
-
-        def sync_models
-          refresh = !finished?
-          records(configuration.fetch('model_table')).find_each(batch_size: BATCH_SIZE) do |model|
-            target = resolve_model(model.id)
-            next unless refresh
-
-            attributes = model.attributes.slice(*target.class.column_names).except(target.class.primary_key)
-            target.assign_attributes(attributes)
-            target.save! if target.changed?
-          end
-        end
-
-        def sync_chat_models(ids = nil)
-          relation = chats.where(VERSION_COLUMN => 1)
-          relation = relation.where(chats.primary_key => ids) if ids
-          models = records(configuration.fetch('model_table')).all.to_h do |model|
-            [model.id, resolve_model(model.id).id]
-          end
-          models.each do |legacy_id, target_id|
-            relation.where(configuration.fetch('model_foreign_key') => legacy_id)
-                    .where('ruby_llm_model_id IS NULL OR ruby_llm_model_id != ?', target_id)
-                    .update_all(ruby_llm_model_id: target_id)
-          end
-        end
-
-        def verify
-          raise 'The online copy journal is not empty; keep AI activity paused' if changes.exists?
-          raise 'RubyLLM 2.0 backfills are incomplete' unless completed?
-
-          tables = [configuration.fetch('chat_table'), configuration.fetch('message_table'),
-                    'ruby_llm_models', 'ruby_llm_tool_calls', 'ruby_llm_usages']
-          if @connection.adapter_name == 'PostgreSQL'
-            @connection.execute("ANALYZE #{tables.map { |table| qt(table) }.join(', ')}")
-          end
-          Verification.new(connection: @connection, settings: configuration).verify
-        end
-
-        def cleanup
-          table = configuration.fetch('message_table')
-          if @connection.column_exists?(table, :ruby_llm_content)
-            @connection.remove_column(table, :content) if @connection.column_exists?(table, :content)
-            @connection.rename_column(table, :ruby_llm_content, :content)
-          end
-          columns = %i[legacy_key legacy_chat_id]
-          %i[ruby_llm_tool_calls ruby_llm_usages].each do |target|
-            columns.each do |column|
-              next unless @connection.column_exists?(target, column)
-
-              @connection.remove_index(target, column) if @connection.index_exists?(target, column)
-              @connection.remove_column(target, column)
-            end
-          end
-        end
-
-        private
-
-        def acknowledge(events)
-          matches = events.map { |event| "(id = #{Integer(event.id)} AND revision = #{Integer(event.revision)})" }
-          changes.where(matches.join(' OR ')).delete_all if matches.any?
-        end
-
-        def prepare_progress
-          unless @connection.table_exists?(PROGRESS)
-            @connection.create_table(PROGRESS) do |table|
-              table.string :task, null: false, index: { unique: true }
-              table.string :last_id
-              table.boolean :completed, null: false, default: false
-            end
-          end
-          @connection.add_column(PROGRESS, :upper_id, :string) unless @connection.column_exists?(PROGRESS, :upper_id)
-          records(PROGRESS).reset_column_information
-        end
-
-        def copy_batch(task, batch)
-          case task
-          when 'message_content' then copy_content(batch.map(&:id))
-          when 'tool_results' then copy_tools(batch)
-          when 'usages' then upsert(:ruby_llm_usages, usage_attributes(batch))
-          end
-        end
-
-        def copy_content(ids)
-          raw = @connection.column_exists?(messages.table_name, :content_raw) ? q(:content_raw) : 'NULL'
-          text = LegacyContentSQL.new(@connection).render(content: q(:content), raw:)
-          structured = @connection.adapter_name == 'PostgreSQL' ? "#{raw}::jsonb" : raw
-          messages.where(messages.primary_key => ids).update_all(
-            "#{q(:ruby_llm_content)} = #{text}, #{q(:raw_content)} = #{structured}"
-          )
-        end
-
-        def copy_tools(batch)
-          ids = batch.map(&:id)
-          source = records(configuration.fetch('tool_call_table'))
-          key = configuration.fetch('message_foreign_key')
-          result_key = configuration.fetch('tool_call_foreign_key')
-          calls = source.where(key => ids).to_a
-          results = messages.where(result_key => calls.map(&:id)).to_a.group_by { |message| message[result_key] }
-          chat_ids = batch.to_h { |message| [message.id, message[configuration.fetch('chat_foreign_key')]] }
-          target = records(:ruby_llm_tool_calls)
-          attributes = calls.map do |call|
-            tool_attributes(call, results.fetch(call.id, []), chat_ids.fetch(call[key]), target)
-          end
-          verify_tool_ids(target, attributes)
-          upsert(:ruby_llm_tool_calls, attributes)
-        end
-
-        def tool_attributes(call, results, chat_id, target)
-          raise "Multiple messages reference tool call #{call.id}" if results.size > 1
-
-          call.attributes.slice(*target.column_names).except(target.primary_key).symbolize_keys.merge(
-            legacy_key: call.id.to_s, legacy_chat_id: chat_id.to_s,
-            message_type: configuration.fetch('message_class'),
-            message_id: call[configuration.fetch('message_foreign_key')],
-            result_type: results.any? ? configuration.fetch('message_class') : nil, result_id: results.first&.id
-          )
-        end
-
-        def verify_tool_ids(target, attributes)
-          provider_ids = attributes.map { |row| row[:tool_call_id] }
-          reserved = target.where(tool_call_id: provider_ids).pluck(:tool_call_id, :legacy_key).to_h
-          attributes.each do |row|
-            next unless reserved.key?(row[:tool_call_id]) && reserved[row[:tool_call_id]] != row[:legacy_key]
-
-            raise "Tool call #{row[:legacy_key]} repeats provider ID #{row[:tool_call_id].inspect}"
-          end
-        end
-
-        def usage_attributes(batch)
-          models = records(configuration.fetch('model_table')).all.index_by(&:id)
-          chat_ids = batch.map { |message| message[configuration.fetch('chat_foreign_key')] }
-          chat_models = chats.where(chats.primary_key => chat_ids)
-                             .pluck(chats.primary_key, configuration.fetch('model_foreign_key')).to_h
-          batch.filter_map do |message|
-            source = message.attributes
-            tokens = legacy_tokens(source)
-            costs = legacy_costs(source)
-            next unless source['role'] == 'assistant' || (tokens.values + costs.values).any? { |value| !value.nil? }
-
-            chat_id = source.fetch(configuration.fetch('chat_foreign_key'))
-            provider, model_id = usage_identity(message, models, chat_models[chat_id])
-
-            tokens.merge(costs).symbolize_keys.merge(
-              legacy_key: message.id.to_s, chat_type: configuration.fetch('chat_class'), chat_id:,
-              message_type: configuration.fetch('message_class'), message_id: message.id,
-              operation: 'chat', provider:, model: model_id, status: 'succeeded',
-              created_at: source['created_at'], updated_at: source['updated_at']
-            )
-          end
-        end
-
-        def usage_identity(message, models, chat_model_id)
-          key = configuration.fetch('model_foreign_key')
-          model = models[message[key]] || models[chat_model_id]
-          if string_model_reference?(message, key)
-            provider = message.attributes['provider'] || model&.provider
-            raise "Message #{message.id} has no identifiable provider" unless provider
-
-            return [provider, message[key]]
-          end
-          raise "Message #{message.id} has no identifiable legacy model" unless model
-
-          [model.provider, model.model_id]
-        end
-
-        def upsert(table, attributes)
-          return if attributes.empty?
-
-          options = { record_timestamps: false }
-          options[:unique_by] = :legacy_key unless @connection.adapter_name == 'Mysql2'
-          records(table).upsert_all(attributes, **options)
-        end
-
-        def string_model_reference?(message, key)
-          messages.columns_hash[key]&.type == :string && message[key]
-        end
-
-        def affected_chats(events)
-          ids = events.select { |event| event.kind == 'chat' }.map(&:record_id)
-          model_ids = events.select { |event| event.kind == 'model' }.map(&:record_id)
-          if model_ids.any?
-            sync_models
-            key = configuration.fetch('model_foreign_key')
-            ids.concat(chats.where(key => model_ids).pluck(chats.primary_key))
-            ids.concat(messages.where(key => model_ids).distinct.pluck(configuration.fetch('chat_foreign_key')))
-          end
-          ids.map(&:to_s).uniq
-        end
-
-        def reconcile_chats(ids)
-          protected_ids = chats.where(chats.primary_key => ids,
-                                      VERSION_COLUMN => 2).pluck(chats.primary_key).map(&:to_s)
-          ids -= protected_ids
-          return if ids.empty?
-
-          remove_deleted_copies(ids)
-          relation = legacy_messages.where(configuration.fetch('chat_foreign_key') => ids)
-          upper = relation.order(messages.primary_key => :desc).pick(messages.primary_key)
-          relation = relation.where("#{q(messages.primary_key)} <= ?", upper) if upper
-          each_batch(relation) do |batch|
-            messages.transaction do
-              TASKS.each { |task| copy_batch(task, batch) }
-              eligible = usage_attributes(batch).map { |row| row[:legacy_key] }
-              records(:ruby_llm_usages).where(legacy_key: batch.map { |message| message.id.to_s })
-                                       .where.not(legacy_key: eligible).delete_all
-            end
-          end
-          remove_deleted_copies(ids)
-        end
-
-        def remove_deleted_copies(ids)
-          current = messages.select(messages.primary_key)
-          records(:ruby_llm_usages).where(chat_type: configuration.fetch('chat_class'), chat_id: ids)
-                                   .where.not(legacy_key: nil).where.not(message_id: current).delete_all
-          source = records(configuration.fetch('tool_call_table'))
-          records(:ruby_llm_tool_calls).where(legacy_chat_id: ids).where.not(legacy_key: nil).find_in_batches do |batch|
-            existing = source.where(source.primary_key => batch.map(&:legacy_key)).pluck(source.primary_key).map(&:to_s)
-            removed = batch.reject { |call| existing.include?(call.legacy_key) }.map(&:id)
-            records(:ruby_llm_tool_calls).where(id: removed).delete_all if removed.any?
-          end
-        end
-
-        def each_batch(relation, &)
-          relation.find_in_batches(batch_size: BATCH_SIZE, &)
-        end
-
-        def legacy_messages
-          owned = chats.where(VERSION_COLUMN => 1).select(chats.primary_key)
-          messages.where(configuration.fetch('chat_foreign_key') => owned)
-        end
-
-        def messages = records(configuration.fetch('message_table'))
-        def chats = records(configuration.fetch('chat_table'))
-        def changes = records(Journal::TABLE)
-        def q(value) = @connection.quote_column_name(value)
-        def qt(value) = @connection.quote_table_name(value)
-      end
-    end
-  end
-end

data/lib/generators/ruby_llm/upgrade/online_copy_migration/journal.rb

@@ -1,171 +0,0 @@
-# frozen_string_literal: true
-
-module RubyLLM
-  module Generators
-    class OnlineCopyMigration < UpgradeMigration
-      class Journal # :nodoc: all
-        TABLE = :ruby_llm_v2_changes
-        OPERATIONS = %w[INSERT UPDATE DELETE].freeze
-        MESSAGE_COLUMNS = %w[
-          role content content_raw input_tokens output_tokens cached_tokens cache_creation_tokens
-          cache_read_tokens cache_write_tokens thinking_tokens total_cost cost_details created_at updated_at
-        ].freeze
-
-        def initialize(connection:, settings:)
-          @connection = connection
-          @settings = settings
-        end
-
-        def install
-          unless @connection.table_exists?(TABLE)
-            @connection.create_table(TABLE) do |table|
-              table.string :kind, null: false
-              table.string :record_id, null: false
-              table.bigint :revision, null: false, default: 1
-              table.index %i[kind record_id], unique: true
-            end
-          end
-          sources.each do |kind, table|
-            OPERATIONS.each { |operation| install_trigger(kind, table, operation) }
-          end
-        end
-
-        def remove
-          sources.each do |kind, table|
-            OPERATIONS.each do |operation|
-              name = trigger_name(kind, operation)
-              if postgresql?
-                if @connection.table_exists?(table)
-                  @connection.execute("DROP TRIGGER IF EXISTS #{q(name)} ON #{qt(table)}")
-                end
-                @connection.execute("DROP FUNCTION IF EXISTS #{q(name)}()")
-              else
-                @connection.execute("DROP TRIGGER IF EXISTS #{q(name)}")
-              end
-            end
-          end
-          @connection.drop_table(TABLE) if @connection.table_exists?(TABLE)
-        end
-
-        private
-
-        def sources
-          %w[chat message tool_call model].to_h { |kind| [kind, @settings.fetch("#{kind}_table")] }
-        end
-
-        def install_trigger(kind, table, operation)
-          name = trigger_name(kind, operation)
-          return if trigger_exists?(name)
-
-          body = event_sql(kind, operation)
-          condition = operation == 'UPDATE' ? changed_sql(kind, table) : '1 = 1'
-          state = qt(UpgradeMigration::TABLE)
-          condition = "(#{condition}) AND EXISTS (SELECT 1 FROM #{state} WHERE active_version = 1)"
-          sql = if postgresql?
-                  <<~SQL
-                    CREATE FUNCTION #{q(name)}() RETURNS trigger LANGUAGE plpgsql AS $ruby_llm$
-                    BEGIN
-                      IF #{condition} THEN #{body} END IF;
-                      RETURN NULL;
-                    END $ruby_llm$;
-                    CREATE TRIGGER #{q(name)} AFTER #{operation} ON #{qt(table)}
-                    FOR EACH ROW EXECUTE FUNCTION #{q(name)}();
-                  SQL
-                elsif mysql?
-                  <<~SQL
-                    CREATE TRIGGER #{q(name)} AFTER #{operation} ON #{qt(table)} FOR EACH ROW
-                    BEGIN IF #{condition} THEN #{body} END IF; END
-                  SQL
-                else
-                  <<~SQL
-                    CREATE TRIGGER #{q(name)} AFTER #{operation} ON #{qt(table)} FOR EACH ROW
-                    WHEN #{condition} BEGIN #{body} END
-                  SQL
-                end
-          @connection.execute(sql)
-        end
-
-        def trigger_exists?(name)
-          quoted = @connection.quote(name)
-          sql = if postgresql?
-                  "SELECT 1 FROM pg_trigger WHERE tgname = #{quoted} AND tgrelid IN " \
-                    '(SELECT oid FROM pg_class WHERE relnamespace = current_schema()::regnamespace)'
-                elsif mysql?
-                  "SELECT 1 FROM information_schema.triggers WHERE trigger_name = #{quoted} " \
-                    'AND trigger_schema = DATABASE()'
-                else
-                  "SELECT 1 FROM sqlite_master WHERE type = 'trigger' AND name = #{quoted}"
-                end
-          @connection.select_value(sql).present?
-        end
-
-        def changed_sql(kind, table)
-          columns = @connection.columns(table).map(&:name)
-          columns &= source_columns(kind) unless kind == 'tool_call'
-          columns.map do |column|
-            old = "OLD.#{q(column)}"
-            new = "NEW.#{q(column)}"
-            if postgresql?
-              "#{old}::text IS DISTINCT FROM #{new}::text"
-            elsif mysql?
-              "NOT (CAST(#{old} AS BINARY) <=> CAST(#{new} AS BINARY))"
-            else
-              "#{old} IS NOT #{new}"
-            end
-          end.join(' OR ')
-        end
-
-        def source_columns(kind)
-          primary = @connection.primary_key(sources.fetch(kind))
-          case kind
-          when 'chat'
-            [primary, @settings.fetch('model_foreign_key')]
-          when 'message'
-            [primary,
-             *@settings.values_at('chat_foreign_key', 'model_foreign_key', 'tool_call_foreign_key')] + MESSAGE_COLUMNS
-          when 'model'
-            [primary, 'provider', 'model_id']
-          end
-        end
-
-        def event_sql(kind, operation)
-          images = if operation == 'UPDATE'
-                     %w[OLD NEW]
-                   else
-                     [operation == 'INSERT' ? 'NEW' : 'OLD']
-                   end
-          images.map do |row|
-            id = event_owner(kind, row)
-            event = kind == 'model' ? 'model' : 'chat'
-            "INSERT INTO #{qt(TABLE)} (kind, record_id) SELECT '#{event}', #{id} WHERE #{id} IS NOT NULL #{upsert_sql};"
-          end.join("\n")
-        end
-
-        def event_owner(kind, row)
-          case kind
-          when 'message'
-            "#{row}.#{q(@settings.fetch('chat_foreign_key'))}"
-          when 'tool_call'
-            "(SELECT #{q(@settings.fetch('chat_foreign_key'))} FROM #{qt(sources.fetch('message'))} " \
-            "WHERE #{q(@connection.primary_key(sources.fetch('message')))} = " \
-            "#{row}.#{q(@settings.fetch('message_foreign_key'))})"
-          else
-            "#{row}.#{q(@connection.primary_key(sources.fetch(kind)))}"
-          end
-        end
-
-        def upsert_sql
-          return 'ON DUPLICATE KEY UPDATE revision = revision + 1' if mysql?
-
-          "ON CONFLICT (kind, record_id) DO UPDATE SET revision = #{qt(TABLE)}.revision + 1"
-        end
-
-        def trigger_name(kind, operation) = "ruby_llm_v2_#{kind}_#{operation.downcase}"
-        def postgresql? = @connection.adapter_name == 'PostgreSQL'
-        def mysql? = @connection.adapter_name == 'Mysql2'
-        def q(value) = @connection.quote_column_name(value)
-        def qt(value) = @connection.quote_table_name(value)
-      end
-    end
-  end
-end