ruby_llm 1.16.0 → 2.0.0.rc4

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

@@ -0,0 +1,197 @@
+# frozen_string_literal: true
+
+require_relative '../legacy_content_sql'
+
+module RubyLLM
+  module Generators
+    class OnlineCopyMigration < UpgradeMigration
+      class Verification # :nodoc: all
+        def initialize(connection:, settings:)
+          @connection = connection
+          @settings = settings
+        end
+
+        def verify
+          verify_chats
+          verify_content
+          verify_usages
+          verify_tools
+        end
+
+        private
+
+        def verify_chats
+          mismatch = %w[provider model_id].map do |column|
+            different("copied.#{q(column)}", "original.#{q(column)}", string: true)
+          end
+          reject_mismatch('chat model', <<~SQL)
+            SELECT c.#{pk('chat')} FROM #{table('chat')} c
+            LEFT JOIN #{table('model')} original ON original.#{pk('model')} = c.#{q(@settings.fetch('model_foreign_key'))}
+            LEFT JOIN #{qt(:ruby_llm_models)} copied ON copied.#{pk('model')} = c.ruby_llm_model_id
+            WHERE c.ruby_llm_version = 1
+              AND (original.#{pk('model')} IS NULL OR copied.#{pk('model')} IS NULL OR #{mismatch.join(' OR ')})
+          SQL
+        end
+
+        def verify_content
+          raw = source(:content_raw)
+          rendered = LegacyContentSQL.new(@connection).render(content: "m.#{q(:content)}", raw:)
+          fields = { ruby_llm_content: rendered, raw_content: raw }
+          failed = fields.map do |column, expected|
+            different("m.#{q(column)}", expected, json: column == :raw_content, string: true)
+          end
+          reject_mismatch('message content', "SELECT m.#{pk('message')} #{message_from} " \
+                                             "WHERE c.ruby_llm_version = 1 AND (#{failed.join(' OR ')})")
+        end
+
+        def verify_usages
+          fields = usage_fields
+          eligible = fields.slice(*%i[input_tokens output_tokens cache_read_tokens cache_write_tokens thinking_tokens
+                                      input_cost output_cost cache_read_cost cache_write_cost thinking_cost total_cost])
+                           .values.map { |value| "#{value} IS NOT NULL" }
+          eligible << "m.#{q(:role)} = 'assistant'"
+          strings = %i[chat_type message_type provider model operation status]
+          failed = fields.map do |column, expected|
+            different("u.#{q(column)}", expected, string: strings.include?(column))
+          end
+          reject_mismatch('usage entry', <<~SQL)
+            SELECT m.#{pk('message')} #{message_from}
+            LEFT JOIN #{table('model')} cm ON cm.#{pk('model')} = c.#{q(@settings.fetch('model_foreign_key'))}
+            #{message_model_join}
+            LEFT JOIN #{qt(:ruby_llm_usages)} u
+              ON u.legacy_key = #{copied_key(:ruby_llm_usages, "m.#{pk('message')}")}
+            WHERE c.ruby_llm_version = 1 AND (#{eligible.join(' OR ')})
+              AND (u.id IS NULL OR #{failed.join(' OR ')})
+          SQL
+        end
+
+        def usage_fields
+          fields = {
+            input_tokens: source(:input_tokens), output_tokens: source(:output_tokens),
+            cache_read_tokens: coalesce(source(:cache_read_tokens), source(:cached_tokens)),
+            cache_write_tokens: coalesce(source(:cache_write_tokens), source(:cache_creation_tokens)),
+            thinking_tokens: source(:thinking_tokens)
+          }
+          %w[input output cache_read cache_write thinking].each { |key| fields[:"#{key}_cost"] = cost(key) }
+          fields[:total_cost] = coalesce(source(:total_cost), cost('total'))
+          provider, model = model_identity
+          fields.merge(chat_type: value(@settings.fetch('chat_class')), chat_id: "c.#{pk('chat')}",
+                       message_type: value(@settings.fetch('message_class')), message_id: "m.#{pk('message')}",
+                       provider:, model:, operation: "'chat'", status: "'succeeded'")
+        end
+
+        def model_identity
+          key = @settings.fetch('model_foreign_key')
+          if message_columns[key]&.type == :string
+            ["COALESCE(#{source(:provider)}, cm.provider)", "COALESCE(#{source(key)}, cm.model_id)"]
+          elsif message_columns.key?(key)
+            ['COALESCE(mm.provider, cm.provider)', 'COALESCE(mm.model_id, cm.model_id)']
+          else
+            ['cm.provider', 'cm.model_id']
+          end
+        end
+
+        def message_model_join
+          key = @settings.fetch('model_foreign_key')
+          return '' if !message_columns.key?(key) || message_columns[key].type == :string
+
+          "LEFT JOIN #{table('model')} mm ON mm.#{pk('model')} = m.#{q(key)}"
+        end
+
+        def cost(key)
+          return 'NULL' unless message_columns.key?('cost_details')
+
+          details = source(:cost_details)
+          case @connection.adapter_name
+          when 'PostgreSQL' then "NULLIF(#{details}::jsonb ->> '#{key}', '')::numeric"
+          when 'Mysql2'
+            "CAST(NULLIF(JSON_UNQUOTE(JSON_EXTRACT(#{details}, '$.#{key}')), 'null') AS DECIMAL(16, 10))"
+          else "json_extract(#{details}, '$.#{key}')"
+          end
+        end
+
+        def verify_tools
+          result_type = "CASE WHEN r.#{pk('message')} IS NOT NULL THEN #{value(@settings.fetch('message_class'))} END"
+          fields = { message_id: "s.#{q(@settings.fetch('message_foreign_key'))}",
+                     message_type: value(@settings.fetch('message_class')), result_id: "r.#{pk('message')}",
+                     result_type:,
+                     tool_call_id: 's.tool_call_id', name: 's.name', arguments: 's.arguments' }
+          strings = %i[message_type result_type tool_call_id name arguments]
+          failed = fields.map do |column, expected|
+            different("t.#{q(column)}", expected, json: column == :arguments, string: strings.include?(column))
+          end
+          reject_mismatch('tool call', <<~SQL)
+            SELECT s.#{pk('tool_call')} FROM #{table('tool_call')} s
+            LEFT JOIN #{table('message')} m ON m.#{pk('message')} = s.#{q(@settings.fetch('message_foreign_key'))}
+            LEFT JOIN #{table('chat')} c ON c.#{pk('chat')} = m.#{q(@settings.fetch('chat_foreign_key'))}
+            LEFT JOIN #{table('message')} r ON r.#{q(@settings.fetch('tool_call_foreign_key'))} = s.#{pk('tool_call')}
+            LEFT JOIN #{qt(:ruby_llm_tool_calls)} t
+              ON t.legacy_key = #{copied_key(:ruby_llm_tool_calls, "s.#{pk('tool_call')}")}
+            WHERE c.ruby_llm_version = 1 AND (t.id IS NULL OR #{failed.join(' OR ')})
+          SQL
+        end
+
+        def message_from
+          "FROM #{table('message')} m JOIN #{table('chat')} c " \
+            "ON c.#{pk('chat')} = m.#{q(@settings.fetch('chat_foreign_key'))}"
+        end
+
+        def different(left, right, json: false, string: false)
+          left, right = comparable([left, right], json:, string:)
+          case @connection.adapter_name
+          when 'PostgreSQL' then "#{left} IS DISTINCT FROM #{right}"
+          when 'Mysql2' then "NOT (#{left} <=> #{right})"
+          else "#{left} IS NOT #{right}"
+          end
+        end
+
+        def comparable(expressions, json:, string:)
+          case @connection.adapter_name
+          when 'PostgreSQL'
+            json ? expressions.map { |expression| "(#{expression})::jsonb" } : expressions
+          when 'Mysql2'
+            string ? expressions.map { |expression| "CAST(#{expression} AS BINARY)" } : expressions
+          else
+            json ? expressions.map { |expression| "json(#{expression})" } : expressions
+          end
+        end
+
+        def reject_mismatch(label, sql)
+          id = @connection.select_value("#{sql.strip} LIMIT 1")
+          raise "Record #{id} did not preserve its #{label}" if id
+        end
+
+        def text(expression) = mysql? ? "CAST(#{expression} AS BINARY)" : "CAST(#{expression} AS TEXT)"
+
+        def copied_key(table, expression)
+          return text(expression) unless mysql?
+
+          collation = @connection.columns(table).find { |column| column.name == 'legacy_key' }.collation
+          charset = collation.split('_').first
+          "CAST(#{expression} AS CHAR CHARACTER SET #{q(charset)}) COLLATE #{q(collation)}"
+        end
+
+        def coalesce(*expressions)
+          values = expressions.reject { |expression| expression == 'NULL' }
+          return 'NULL' if values.empty?
+          return values.first if values.one?
+
+          "COALESCE(#{values.join(', ')})"
+        end
+
+        def source(column) = message_columns.key?(column.to_s) ? "m.#{q(column)}" : 'NULL'
+
+        def message_columns
+          @message_columns ||= @connection.columns(@settings.fetch('message_table')).index_by(&:name)
+        end
+
+        def mysql? = @connection.adapter_name == 'Mysql2'
+        def table(kind) = qt(@settings.fetch("#{kind}_table"))
+        def pk(kind) = q(@connection.primary_key(@settings.fetch("#{kind}_table")))
+        def value(value) = @connection.quote(value)
+        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.rb

@@ -0,0 +1,174 @@
+# frozen_string_literal: true
+
+require_relative 'upgrade_migration'
+require_relative 'online_copy_migration/journal'
+require_relative 'online_copy_migration/data'
+
+module RubyLLM
+  module Generators
+    class OnlineCopyMigration < UpgradeMigration # :nodoc: all
+      def online? = true
+
+      def prepare(settings)
+        with_migration_lock do
+          super(settings.merge(online: true), &nil)
+          unless @connection.column_exists?(TABLE, :epoch)
+            @connection.add_column(TABLE, :epoch, :bigint, null: false, default: 0)
+          end
+          yield if block_given?
+          data.prepare
+          journal.install
+        end
+      end
+
+      def verify_settings(settings)
+        super(settings.merge(online: true))
+      end
+
+      def copy_table(source, target)
+        create_copy_table(source, target) unless @connection.table_exists?(target)
+      end
+
+      def copy_model_references
+        nil
+      end
+
+      def backfill
+        with_migration_lock do
+          state = states.first!
+          unless state.active_version == 1 && %w[preparing active].include?(state.status)
+            raise 'Online copy backfill requires RubyLLM 1.16 to be active'
+          end
+          raise 'Use resume after rolling back a finished upgrade' if data.finished?
+
+          data.backfill
+          data.catch_up(passes: 2)
+          state.update!(needs_reconcile: false)
+        end
+      end
+
+      def finish(discard_incomplete_tool_calls: false)
+        with_migration_lock do
+          state = states.first!
+          return if state.active_version == 2 && state.status == 'active' && data.finished?
+
+          raise 'Run the copy backfill before finish' unless data.completed?
+
+          pause(state, from: 1, statuses: %w[preparing active finishing])
+          if discard_incomplete_tool_calls
+            discarded = data.discard_incomplete_tool_calls
+            ::ActiveRecord::Migration.say "Discarded #{discarded} incomplete legacy tool calls"
+          end
+          data.catch_up
+          data.verify
+          verify_completed_work
+          yield
+          journal.install
+          state.update!(active_version: 2, status: 'active', needs_reconcile: false)
+        end
+      end
+
+      def rollback
+        with_migration_lock do
+          state = states.first!
+          unless data.finished?
+            raise 'The unfinished upgrade is not running on 1.16' unless state.active_version == 1
+
+            state.update!(status: 'preparing', needs_reconcile: true)
+            return
+          end
+          state.with_lock do
+            require_active_version(state, 2)
+            verify_completed_work
+            state.update!(active_version: 1, epoch: state.epoch + 1)
+          end
+        end
+      end
+
+      def resume
+        with_migration_lock do
+          raise 'Finish the initial copy migrations before resuming 2.0' unless data.finished?
+
+          state = states.first!
+          pause(state, from: 1, statuses: %w[active finishing])
+          data.catch_up
+          data.verify
+          verify_completed_work
+          state.update!(active_version: 2, status: 'active', needs_reconcile: false)
+        end
+      end
+
+      def finalize
+        with_migration_lock { super }
+      end
+
+      def cleanup
+        with_migration_lock do
+          verify_cleanup
+          journal.remove
+          yield if block_given?
+          data.cleanup
+          super
+        end
+      end
+
+      private
+
+      def pause(state, from:, statuses:)
+        state.with_lock do
+          unless state.active_version == from && statuses.include?(state.status)
+            raise 'The copy upgrade is not ready for this version switch'
+          end
+
+          state.update!(status: 'finishing', epoch: state.epoch + 1) unless state.status == 'finishing'
+        end
+      end
+
+      def data
+        @data ||= Data.new(connection: @connection, settings: configuration)
+      end
+
+      def journal
+        @journal ||= Journal.new(connection: @connection, settings: configuration)
+      end
+
+      def with_migration_lock(&)
+        return with_mysql_lock(&) if @connection.adapter_name == 'Mysql2'
+        return with_sqlite_lock(&) if @connection.adapter_name == 'SQLite'
+
+        key = "hashtextextended(current_database() || '.' || current_schema() || '.ruby_llm_upgrade', 0)"
+        locked = @connection.select_value("SELECT pg_try_advisory_lock(#{key})")
+        raise 'Another RubyLLM copy migration is running' unless locked
+
+        begin
+          yield
+        ensure
+          @connection.execute("SELECT pg_advisory_unlock(#{key})")
+        end
+      end
+
+      def with_mysql_lock
+        key = @connection.quote("ruby_llm_upgrade:#{@connection.current_database}".slice(0, 64))
+        locked = @connection.select_value("SELECT GET_LOCK(#{key}, 0)")
+        raise 'Another RubyLLM copy migration is running' unless locked == 1
+
+        begin
+          yield
+        ensure
+          @connection.execute("SELECT RELEASE_LOCK(#{key})")
+        end
+      end
+
+      def with_sqlite_lock
+        database = @connection.pool.db_config.database
+        return yield if database == ':memory:'
+
+        File.open("#{database}.ruby_llm_upgrade.lock", File::RDWR | File::CREAT, 0o600) do |file|
+          raise 'Another RubyLLM copy migration is running' unless file.flock(File::LOCK_EX | File::LOCK_NB)
+
+          yield
+        end
+      end
+    end
+  end
+end