Class: RubyLLM::Generators::OnlineCopyMigration::Data

Inherits:
UpgradeMigration show all
Defined in:
lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb

Overview

:nodoc: all

Constant Summary collapse

BATCH_SIZE = 
10_000
PROGRESS = 
:ruby_llm_v2_backfills
TASKS = 
%w[message_content tool_results usages].freeze

Constants inherited from UpgradeMigration

UpgradeMigration::TABLE, UpgradeMigration::VERSION_COLUMN

Instance Method Summary collapse

Methods inherited from UpgradeMigration

#activate, #copy_model_references, #copy_table, #finalize, for, #online?, #resume, #rollback, #verify_cleanup, #verify_settings

Constructor Details

#initialize(connection:, settings:) ⇒ Data

Returns a new instance of Data.



15
16
17
18
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 15

def initialize(connection:, settings:)
  super(connection:)
  @configuration = settings
end

Instance Method Details

#backfill ⇒ Object 



37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 37

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

#catch_up(passes: nil) ⇒ Object 



75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 75

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

#cleanup ⇒ Object 



140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 140

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

#completed? ⇒ Boolean

Returns:

  • (Boolean) 


94
95
96
97
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 94

def completed?
  @connection.table_exists?(PROGRESS) &&
    (TASKS - records(PROGRESS).where(completed: true).pluck(:task)).empty?
end

#discard_incomplete_tool_calls ⇒ Object 



66
67
68
69
70
71
72
73
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 66

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

#finished? ⇒ Boolean

Returns:

  • (Boolean) 


99
100
101
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 99

def finished?
  @connection.table_exists?(PROGRESS) && records(PROGRESS).where(task: 'finished', completed: true).exists?
end

#prepare ⇒ Object 



20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 20

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

#sync_chat_models(ids = nil) ⇒ Object 



115
116
117
118
119
120
121
122
123
124
125
126
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 115

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

#sync_models ⇒ Object 



103
104
105
106
107
108
109
110
111
112
113
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 103

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

#verify ⇒ Object 



128
129
130
131
132
133
134
135
136
137
138
 
# File 'lib/generators/ruby_llm/upgrade/online_copy_migration/data.rb', line 128

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