ruby_llm 1.16.0 → 2.0.0.rc4

data/lib/generators/ruby_llm/upgrade/templates/finish_v2_upgrade.rb.tt

@@ -0,0 +1,247 @@
+# frozen_string_literal: true
+<% if copy_mode? -%>
+
+require 'generators/ruby_llm/upgrade/upgrade_migration'
+<% end -%>
+
+class <%= finish_migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
+  disable_ddl_transaction!
+
+  PROGRESS_TABLE = :ruby_llm_v2_backfills
+  BACKFILL_TASKS = %w[message_content tool_results usages].freeze
+
+  def up
+<% if copy_mode? -%>
+    RubyLLM::Generators::UpgradeMigration.for.verify_settings(<%= copy_upgrade_settings.inspect %>)
+    RubyLLM::Generators::UpgradeMigration.for.finish<% if discard_incomplete_tool_calls? %>(discard_incomplete_tool_calls: true)<% end %> do
+      verify_completed_backfills
+      enforce_required_defaults
+      relax_legacy_message_constraints
+      mark_finished
+    end
+<% else -%>
+    raise 'Generate this migration with --mode copy' if table_exists?(:ruby_llm_v2_upgrades)
+    verify_completed_backfills
+<% if postgresql? -%>
+    tables = %i[<%= chat_table_name %> <%= message_table_name %> ruby_llm_models ruby_llm_tool_calls ruby_llm_usages]
+    with_upgrade_safety { execute "ANALYZE #{tables.map { |table| quote_table(table) }.join(', ')}" }
+<% end -%>
+    verify_message_content
+    verify_tool_results
+    verify_usages
+    enforce_required_defaults
+    relax_legacy_message_constraints
+    mark_finished
+<% end -%>
+  end
+
+  def down
+    raise ActiveRecord::IrreversibleMigration, 'RubyLLM 2.0 changes persisted record ownership'
+  end
+
+  private
+
+  def relax_legacy_message_constraints
+<% if copy_mode? -%>
+    model_column = :<%= v1_model_foreign_key %>
+    if column_definition(:<%= chat_table_name %>, model_column)&.null == false
+      with_upgrade_safety { change_column_null :<%= chat_table_name %>, model_column, true }
+    end
+    connection.foreign_keys(:<%= v1_tool_call_table_name %>).each do |key|
+      next unless key.column == '<%= message_foreign_key %>'
+
+      remove_foreign_key :<%= v1_tool_call_table_name %>, column: key.column
+    end
+<% end -%>
+    %i[<%= v1_model_foreign_key %> <%= v1_tool_call_foreign_key %>].each do |column|
+      next unless column_definition(:<%= message_table_name %>, column)&.null == false
+
+      with_upgrade_safety { change_column_null :<%= message_table_name %>, column, true }
+    end
+  end
+
+  def mark_finished
+    progress = migration_record(PROGRESS_TABLE)
+    return if progress.where(task: 'finished', completed: true).exists?
+
+    progress.insert_all!([{ task: 'finished', completed: true }], returning: false)
+  end
+
+  def enforce_required_defaults
+    {
+      <%= chat_table_name %>: {cancelled: false},
+      <%= message_table_name %>: {cache_until_here: false},
+      ruby_llm_tool_calls: {message_type: '<%= message_model_name %>'}
+    }.each do |table, attributes|
+      migration_record(table).where(attributes.transform_values { nil }).in_batches(of: 10_000) do |batch|
+        batch.update_all(attributes)
+      end
+    end
+<% if copy_mode? -%>
+    enforce_not_null(:<%= chat_table_name %>, :ruby_llm_model_id)
+<% end -%>
+    enforce_not_null(:<%= chat_table_name %>, :cancelled)
+    enforce_not_null(:<%= message_table_name %>, :cache_until_here)
+    enforce_not_null(:ruby_llm_tool_calls, :message_type)
+  end
+
+  def enforce_not_null(table, column)
+    missing = migration_record(table).where(column => nil).limit(1).pick(connection.primary_key(table))
+    raise "#{table} id #{missing} has a NULL #{column}" if missing
+    return unless column_definition(table, column)&.null
+
+<% if postgresql? -%>
+    constraint = "chk_ruby_llm_v2_#{column}_not_null"
+    unless check_constraint_exists?(table, name: constraint)
+      add_check_constraint table, "#{quote_column(column)} IS NOT NULL", name: constraint, validate: false
+    end
+    validate_check_constraint table, name: constraint
+    with_upgrade_safety { change_column_null table, column, false }
+    remove_check_constraint table, name: constraint
+<% else -%>
+    with_upgrade_safety { change_column_null table, column, false }
+<% end -%>
+  end
+
+  def verify_completed_backfills
+    raise "Run <%= backfill_migration_class_name %> before this migration" unless table_exists?(PROGRESS_TABLE)
+
+    completed = migration_record(PROGRESS_TABLE).where(completed: true).pluck(:task)
+    missing = BACKFILL_TASKS - completed
+    return if missing.empty?
+
+    raise "RubyLLM 2.0 backfills are incomplete: #{missing.join(', ')}"
+  end
+
+<% unless copy_mode? -%>
+  def verify_message_content
+    return unless column_exists?(:<%= message_table_name %>, :content_raw)
+
+    missing = migration_record(:<%= message_table_name %>)
+              .where.not(content_raw: nil)
+              .where('raw_content IS NULL OR content IS NULL')
+              .limit(1)
+              .pick(connection.primary_key(:<%= message_table_name %>))
+    raise "Message #{missing} did not preserve content_raw" if missing
+  end
+
+  def verify_tool_results
+    result_reference = :<%= v1_tool_call_foreign_key %>
+    return unless column_exists?(:<%= message_table_name %>, result_reference)
+
+    missing = select_value(<<~SQL)
+      SELECT legacy_messages.#{quoted_primary_key(:<%= message_table_name %>)}
+        FROM #{quote_table(:<%= message_table_name %>)} legacy_messages
+<% if copy_mode? -%>
+        LEFT JOIN #{quote_table(:<%= v1_tool_call_table_name %>)} legacy_tool_calls
+          ON #{message_value(result_reference)} = legacy_tool_calls.#{quoted_primary_key(:<%= v1_tool_call_table_name %>)}
+        LEFT JOIN #{quote_table(:ruby_llm_tool_calls)} migrated_tool_calls
+          ON legacy_tool_calls.tool_call_id = migrated_tool_calls.tool_call_id
+         AND legacy_tool_calls.#{quote_column(:<%= message_foreign_key %>)} = migrated_tool_calls.message_id
+<% else -%>
+        LEFT JOIN #{quote_table(:ruby_llm_tool_calls)} migrated_tool_calls
+          ON #{message_value(result_reference)} = migrated_tool_calls.#{quoted_primary_key(:ruby_llm_tool_calls)}
+<% end -%>
+         AND migrated_tool_calls.result_type = #{connection.quote('<%= message_model_name %>')}
+         AND migrated_tool_calls.result_id = legacy_messages.#{quoted_primary_key(:<%= message_table_name %>)}
+       WHERE #{message_value(result_reference)} IS NOT NULL
+         AND migrated_tool_calls.#{quoted_primary_key(:ruby_llm_tool_calls)} IS NULL
+       LIMIT 1
+    SQL
+    raise "Message #{missing} did not preserve its tool result association" if missing
+  end
+
+  def verify_usages
+    joins = usage_identity_joins
+    conditions = legacy_usage_conditions
+    return if conditions.empty?
+
+    missing = select_value(<<~SQL)
+      SELECT #{message_value(connection.primary_key(:<%= message_table_name %>))}
+        FROM #{quote_table(:<%= message_table_name %>)} legacy_messages
+        #{joins}
+       WHERE (#{conditions.join(' OR ')})
+         AND NOT EXISTS (
+           SELECT 1
+             FROM #{quote_table(:ruby_llm_usages)} migrated_usages
+            WHERE migrated_usages.message_type = #{connection.quote('<%= message_model_name %>')}
+              AND migrated_usages.message_id = #{message_value(connection.primary_key(:<%= message_table_name %>))}
+              AND migrated_usages.operation = 'chat'
+              AND migrated_usages.status = 'succeeded'
+         )
+       LIMIT 1
+    SQL
+    raise "Message #{missing} did not receive its RubyLLM usage entry" if missing
+  end
+
+  def usage_identity_joins
+    joins = []
+    registry = quote_table(:ruby_llm_models)
+    registry_key = quoted_primary_key(:ruby_llm_models)
+    if column_exists?(:<%= message_table_name %>, :<%= v1_model_foreign_key %>) &&
+       column_definition(:<%= message_table_name %>, :<%= v1_model_foreign_key %>).type != :string
+      joins << "LEFT JOIN #{registry} legacy_message_models " \
+               "ON #{message_value(:<%= v1_model_foreign_key %>)} = legacy_message_models.#{registry_key}"
+    end
+    if column_exists?(:<%= chat_table_name %>, :ruby_llm_model_id)
+      joins << "LEFT JOIN #{quote_table(:<%= chat_table_name %>)} legacy_chats " \
+               "ON #{message_value(:<%= chat_foreign_key %>)} = legacy_chats.#{quoted_primary_key(:<%= chat_table_name %>)}"
+      joins << "LEFT JOIN #{registry} legacy_chat_models " \
+               "ON legacy_chats.#{quote_column(:ruby_llm_model_id)} = legacy_chat_models.#{registry_key}"
+    end
+    joins.join("\n")
+  end
+
+  def legacy_usage_conditions
+    columns = %i[
+      input_tokens output_tokens cached_tokens cache_creation_tokens cache_read_tokens cache_write_tokens
+      thinking_tokens total_cost
+    ]
+    conditions = columns.filter_map do |column|
+      "#{message_value(column)} IS NOT NULL" if column_exists?(:<%= message_table_name %>, column)
+    end
+    if column_exists?(:<%= message_table_name %>, :cost_details)
+      costs = %w[input output cache_read cache_write thinking total].map do |key|
+        "#{legacy_cost_detail(key)} IS NOT NULL"
+      end
+      conditions << "(#{costs.join(' OR ')})"
+    end
+    conditions << "#{message_value(:role)} = 'assistant'" if column_exists?(:<%= message_table_name %>, :role)
+    conditions
+  end
+
+  def legacy_cost_detail(key)
+    details = message_value(:cost_details)
+<% if postgresql? -%>
+    "NULLIF(#{details}::jsonb ->> '#{key}', '')::numeric"
+<% elsif mysql? -%>
+    "JSON_UNQUOTE(JSON_EXTRACT(#{details}, '$.#{key}'))"
+<% else -%>
+    "json_extract(#{details}, '$.#{key}')"
+<% end -%>
+  end
+
+<% end -%>
+  def with_upgrade_safety(&)
+    return safety_assured(&) if respond_to?(:safety_assured, true)
+
+    yield
+  end
+
+  def column_definition(table, column) = connection.columns(table).find { |candidate| candidate.name == column.to_s }
+<% unless copy_mode? -%>
+  def message_value(column) = "legacy_messages.#{quote_column(column)}"
+  def quoted_primary_key(table) = quote_column(connection.primary_key(table))
+  def quote_table(table) = connection.quote_table_name(table)
+<% end -%>
+<% if postgresql? || !copy_mode? -%>
+  def quote_column(column) = connection.quote_column_name(column)
+<% end -%>
+
+  def migration_record(table)
+    Class.new(ActiveRecord::Base) do
+      self.table_name = table.to_s
+      self.inheritance_column = :_type_disabled
+    end
+  end
+end