ruby_llm_swarm 1.9.1

data/lib/generators/ruby_llm/chat_ui/templates/views/models/_model.html.erb.tt

@@ -0,0 +1,16 @@
+<tr id="<%%= dom_id <%= model_model_name.demodulize.underscore %> %>">
+  <td><%%= <%= model_model_name.demodulize.underscore %>.provider %></td>
+  <td><%%= <%= model_model_name.demodulize.underscore %>.name %></td>
+  <td><%%= number_with_delimiter(<%= model_model_name.demodulize.underscore %>.context_window) if <%= model_model_name.demodulize.underscore %>.context_window %></td>
+  <td>
+    <%% if <%= model_model_name.demodulize.underscore %>.pricing && <%= model_model_name.demodulize.underscore %>.pricing['text_tokens'] && <%= model_model_name.demodulize.underscore %>.pricing['text_tokens']['standard'] %>
+      <%% input = <%= model_model_name.demodulize.underscore %>.pricing['text_tokens']['standard']['input_per_million'] %>
+      <%% output = <%= model_model_name.demodulize.underscore %>.pricing['text_tokens']['standard']['output_per_million'] %>
+      <%% if input && output %>
+        $<%%= "%.2f" % input %> / $<%%= "%.2f" % output %>
+      <%% end %>
+    <%% end %>
+  </td>
+  <td><%%= link_to "Show", <%= model_model_name.demodulize.underscore %> %></td>
+  <td><%%= link_to "Start <%= chat_table_name.singularize.humanize.downcase %>", new_<%= chat_variable_name %>_path(model: <%= model_model_name.demodulize.underscore %>.model_id) %></td>
+</tr>

data/lib/generators/ruby_llm/chat_ui/templates/views/models/index.html.erb.tt

@@ -0,0 +1,28 @@
+<p style="color: green"><%%= notice %></p>
+
+<%% content_for :title, "<%= model_model_name.pluralize %>" %>
+
+<h1><%= model_model_name.pluralize %></h1>
+
+<p>
+  <%%= button_to "Refresh <%= model_model_name.pluralize %>", refresh_<%= model_table_name %>_path, method: :post %>
+</p>
+
+<div id="<%= model_table_name %>">
+  <table>
+    <thead>
+      <tr>
+        <th>Provider</th>
+        <th>Model</th>
+        <th>Context Window</th>
+        <th>$/1M tokens (In/Out)</th>
+        <th colspan="2"></th>
+      </tr>
+    </thead>
+    <tbody>
+      <%%= render @<%= model_variable_name.pluralize %> %>
+    </tbody>
+  </table>
+</div>
+
+<%%= link_to "Back to <%= chat_table_name.humanize.downcase %>", <%= chat_table_name %>_path %>

data/lib/generators/ruby_llm/chat_ui/templates/views/models/show.html.erb.tt

@@ -0,0 +1,18 @@
+<%% content_for :title, @<%= model_variable_name %>.name %>
+
+<h1><%%= @<%= model_variable_name %>.name %></h1>
+
+<p><strong>ID:</strong> <%%= @<%= model_variable_name %>.model_id %></p>
+<p><strong>Provider:</strong> <%%= @<%= model_variable_name %>.provider %></p>
+<p><strong>Context Window:</strong> <%%= number_with_delimiter(@<%= model_variable_name %>.context_window) %> tokens</p>
+<p><strong>Max Output:</strong> <%%= number_with_delimiter(@<%= model_variable_name %>.max_output_tokens) %> tokens</p>
+
+<%% if @<%= model_variable_name %>.capabilities.any? %>
+  <p><strong>Capabilities:</strong> <%%= @<%= model_variable_name %>.capabilities.join(", ") %></p>
+<%% end %>
+
+<p>
+  <%%= link_to "Start chat with this model", new_<%= chat_variable_name %>_path(model: @<%= model_variable_name %>.model_id) %> |
+  <%%= link_to "All models", <%= model_table_name %>_path %> |
+  <%%= link_to "Back to chats", <%= chat_table_name %>_path %>
+</p>

data/lib/generators/ruby_llm/generator_helpers.rb

@@ -0,0 +1,194 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Generators
+    # Shared helpers for RubyLLM generators
+    module GeneratorHelpers
+      def parse_model_mappings
+        @model_names = {
+          chat: 'Chat',
+          message: 'Message',
+          tool_call: 'ToolCall',
+          model: 'Model'
+        }
+
+        model_mappings.each do |mapping|
+          if mapping.include?(':')
+            key, value = mapping.split(':', 2)
+            @model_names[key.to_sym] = value.classify
+          end
+        end
+
+        @model_names
+      end
+
+      %i[chat message tool_call model].each do |type|
+        define_method("#{type}_model_name") do
+          @model_names ||= parse_model_mappings
+          @model_names[type]
+        end
+
+        define_method("#{type}_table_name") do
+          table_name_for(send("#{type}_model_name"))
+        end
+
+        define_method("#{type}_variable_name") do
+          variable_name_for(send("#{type}_model_name"))
+        end
+
+        define_method("#{type}_controller_class_name") do
+          controller_class_name_for(send("#{type}_model_name"))
+        end
+
+        define_method("#{type}_job_class_name") do
+          "#{variable_name_for(send("#{type}_model_name")).camelize}ResponseJob"
+        end
+
+        define_method("#{type}_partial") do
+          partial_path_for(send("#{type}_model_name"))
+        end
+      end
+
+      def acts_as_chat_declaration
+        params = []
+
+        add_association_params(params, :messages, message_table_name, message_model_name,
+                               owner_table: chat_table_name, plural: true)
+        add_association_params(params, :model, model_table_name, model_model_name,
+                               owner_table: chat_table_name)
+
+        "acts_as_chat#{" #{params.join(', ')}" if params.any?}"
+      end
+
+      def acts_as_message_declaration
+        params = []
+
+        add_association_params(params, :chat, chat_table_name, chat_model_name,
+                               owner_table: message_table_name)
+        add_association_params(params, :tool_calls, tool_call_table_name, tool_call_model_name,
+                               owner_table: message_table_name, plural: true)
+        add_association_params(params, :model, model_table_name, model_model_name,
+                               owner_table: message_table_name)
+
+        "acts_as_message#{" #{params.join(', ')}" if params.any?}"
+      end
+
+      def acts_as_model_declaration
+        params = []
+
+        add_association_params(params, :chats, chat_table_name, chat_model_name,
+                               owner_table: model_table_name, plural: true)
+
+        "acts_as_model#{" #{params.join(', ')}" if params.any?}"
+      end
+
+      def acts_as_tool_call_declaration
+        params = []
+
+        add_association_params(params, :message, message_table_name, message_model_name,
+                               owner_table: tool_call_table_name)
+
+        "acts_as_tool_call#{" #{params.join(', ')}" if params.any?}"
+      end
+
+      def create_namespace_modules
+        namespaces = []
+
+        [chat_model_name, message_model_name, tool_call_model_name, model_model_name].each do |model_name|
+          if model_name.include?('::')
+            namespace = model_name.split('::').first
+            namespaces << namespace unless namespaces.include?(namespace)
+          end
+        end
+
+        namespaces.each do |namespace|
+          module_path = "app/models/#{namespace.underscore}.rb"
+          next if File.exist?(Rails.root.join(module_path))
+
+          create_file module_path do
+            <<~RUBY
+              module #{namespace}
+                def self.table_name_prefix
+                  "#{namespace.underscore}_"
+                end
+              end
+            RUBY
+          end
+        end
+      end
+
+      def migration_version
+        "[#{Rails::VERSION::MAJOR}.#{Rails::VERSION::MINOR}]"
+      end
+
+      def postgresql?
+        ::ActiveRecord::Base.connection.adapter_name.downcase.include?('postgresql')
+      rescue StandardError
+        false
+      end
+
+      def mysql?
+        ::ActiveRecord::Base.connection.adapter_name.downcase.include?('mysql')
+      rescue StandardError
+        false
+      end
+
+      def table_exists?(table_name)
+        ::ActiveRecord::Base.connection.table_exists?(table_name)
+      rescue StandardError
+        false
+      end
+
+      private
+
+      def add_association_params(params, default_assoc, table_name, model_name, owner_table:, plural: false) # rubocop:disable Metrics/ParameterLists
+        assoc = plural ? table_name.to_sym : table_name.singularize.to_sym
+
+        default_foreign_key = "#{default_assoc}_id"
+        # has_many/has_one: foreign key is on the associated table pointing back to owner
+        # belongs_to:       foreign key is on the owner table pointing to associated table
+        foreign_key = if plural || default_assoc.to_s.pluralize == default_assoc.to_s # has_many or has_one
+                        "#{owner_table.singularize}_id"
+                      else # belongs_to
+                        "#{table_name.singularize}_id"
+                      end
+
+        params << "#{default_assoc}: :#{assoc}" if assoc != default_assoc
+        params << "#{default_assoc.to_s.singularize}_class: '#{model_name}'" if model_name != assoc.to_s.classify
+        params << "#{default_assoc}_foreign_key: :#{foreign_key}" if foreign_key != default_foreign_key
+      end
+
+      # Convert namespaced model names to proper table names
+      # e.g., "Assistant::Chat" -> "assistant_chats" (not "assistant/chats")
+      def table_name_for(model_name)
+        model_name.underscore.pluralize.tr('/', '_')
+      end
+
+      # Convert model name to instance variable name
+      # e.g., "LLM::Chat" -> "llm_chat" (not "llm/chat")
+      def variable_name_for(model_name)
+        model_name.underscore.tr('/', '_')
+      end
+
+      # Convert model name to controller class name
+      # For namespaced models, use Rails convention: "Llm::Chat" -> "Llm::ChatsController"
+      # For regular models: "Chat" -> "ChatsController"
+      def controller_class_name_for(model_name)
+        if model_name.include?('::')
+          parts = model_name.split('::')
+          namespace = parts[0..-2].join('::')
+          resource = parts.last.pluralize
+          "#{namespace}::#{resource}Controller"
+        else
+          "#{model_name.pluralize}Controller"
+        end
+      end
+
+      # Convert model name to partial path
+      # e.g., "LLM::Message" -> "llm/message" (not "llm_message")
+      def partial_path_for(model_name)
+        "#{model_name.underscore.pluralize}/#{model_name.demodulize.underscore}"
+      end
+    end
+  end
+end

data/lib/generators/ruby_llm/install/install_generator.rb

@@ -0,0 +1,106 @@
+# frozen_string_literal: true
+
+require 'rails/generators'
+require 'rails/generators/active_record'
+require_relative '../generator_helpers'
+
+module RubyLLM
+  module Generators
+    # Generator for RubyLLM Rails models and migrations
+    class InstallGenerator < Rails::Generators::Base
+      include Rails::Generators::Migration
+      include RubyLLM::Generators::GeneratorHelpers
+
+      namespace 'ruby_llm:install'
+
+      source_root File.expand_path('templates', __dir__)
+
+      argument :model_mappings, type: :array, default: [], banner: 'chat:ChatName message:MessageName ...'
+
+      class_option :skip_active_storage, type: :boolean, default: false,
+                                         desc: 'Skip ActiveStorage installation and attachment setup'
+
+      desc 'Creates models and migrations for RubyLLM Rails integration\n' \
+           'Usage: rails g ruby_llm:install [chat:ChatName] [message:MessageName] ...'
+
+      def self.next_migration_number(dirname)
+        ::ActiveRecord::Generators::Base.next_migration_number(dirname)
+      end
+
+      def create_migration_files
+        migration_template 'create_chats_migration.rb.tt',
+                           "db/migrate/create_#{chat_table_name}.rb"
+
+        sleep 1 # Ensure different timestamp
+        migration_template 'create_messages_migration.rb.tt',
+                           "db/migrate/create_#{message_table_name}.rb"
+
+        sleep 1 # Ensure different timestamp
+        migration_template 'create_tool_calls_migration.rb.tt',
+                           "db/migrate/create_#{tool_call_table_name}.rb"
+
+        sleep 1 # Ensure different timestamp
+        migration_template 'create_models_migration.rb.tt',
+                           "db/migrate/create_#{model_table_name}.rb"
+
+        sleep 1 # Ensure different timestamp
+        migration_template 'add_references_to_chats_tool_calls_and_messages_migration.rb.tt',
+                           'db/migrate/add_references_to_' \
+                           "#{chat_table_name}_#{tool_call_table_name}_and_#{message_table_name}.rb"
+      end
+
+      def create_model_files
+        create_namespace_modules
+
+        template 'chat_model.rb.tt', "app/models/#{chat_model_name.underscore}.rb"
+        template 'message_model.rb.tt', "app/models/#{message_model_name.underscore}.rb"
+        template 'tool_call_model.rb.tt', "app/models/#{tool_call_model_name.underscore}.rb"
+
+        template 'model_model.rb.tt', "app/models/#{model_model_name.underscore}.rb"
+      end
+
+      def create_initializer
+        template 'initializer.rb.tt', 'config/initializers/ruby_llm.rb'
+      end
+
+      def install_active_storage
+        return if options[:skip_active_storage]
+
+        say '  Installing ActiveStorage for file attachments...', :cyan
+        rails_command 'active_storage:install'
+      end
+
+      def show_install_info
+        say "\n  ✅ RubyLLM installed!", :green
+
+        say '  ✅ ActiveStorage configured for file attachments support', :green unless options[:skip_active_storage]
+
+        say "\n  Next steps:", :yellow
+        say '     1. Run: rails db:migrate'
+        say '     2. Set your API keys in config/initializers/ruby_llm.rb'
+
+        say "     3. Start chatting: #{chat_model_name}.create!(model: 'gpt-4.1-nano').ask('Hello!')"
+
+        say "\n  🚀 Model registry is database-backed!", :cyan
+        say '     Models automatically load from the database'
+        say '     Pass model names as strings - RubyLLM handles the rest!'
+        say "     Specify provider when needed: Chat.create!(model: 'gemini-2.5-flash', provider: 'vertexai')"
+
+        if options[:skip_active_storage]
+          say "\n  📎 Note: ActiveStorage was skipped", :yellow
+          say '     File attachments won\'t work without ActiveStorage.'
+          say '     To enable later:'
+          say '       1. Run: rails active_storage:install && rails db:migrate'
+          say "       2. Add to your #{message_model_name} model: has_many_attached :attachments"
+        end
+
+        say "\n  📚 Documentation: https://rubyllm.com", :cyan
+
+        say "\n  ❤️  Love RubyLLM?", :magenta
+        say '     • ⭐ Star on GitHub: https://github.com/crmne/ruby_llm'
+        say '     • 🐦 Follow for updates: https://x.com/paolino'
+        say "\n"
+      end
+    end
+  end
+end

data/lib/generators/ruby_llm/install/templates/add_references_to_chats_tool_calls_and_messages_migration.rb.tt

@@ -0,0 +1,9 @@
+class AddReferencesTo<%= "#{chat_model_name.gsub('::', '').pluralize}#{tool_call_model_name.gsub('::', '').pluralize}And#{message_model_name.gsub('::', '').pluralize}" %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    add_reference :<%= chat_table_name %>, :<%= model_table_name.singularize %>, foreign_key: true
+    add_reference :<%= tool_call_table_name %>, :<%= message_table_name.singularize %>, null: false, foreign_key: true
+    add_reference :<%= message_table_name %>, :<%= chat_table_name.singularize %>, null: false, foreign_key: true
+    add_reference :<%= message_table_name %>, :<%= model_table_name.singularize %>, foreign_key: true
+    add_reference :<%= message_table_name %>, :<%= tool_call_table_name.singularize %>, foreign_key: true
+  end
+end

data/lib/generators/ruby_llm/install/templates/chat_model.rb.tt

@@ -0,0 +1,3 @@
+class <%= chat_model_name %> < ApplicationRecord
+  <%= acts_as_chat_declaration %>
+end

data/lib/generators/ruby_llm/install/templates/create_chats_migration.rb.tt

@@ -0,0 +1,7 @@
+class Create<%= chat_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :<%= chat_table_name %> do |t|
+      t.timestamps
+    end
+  end
+end

data/lib/generators/ruby_llm/install/templates/create_messages_migration.rb.tt

@@ -0,0 +1,16 @@
+class Create<%= message_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :<%= message_table_name %> do |t|
+      t.string :role, null: false
+      t.text :content
+      t.json :content_raw
+      t.integer :input_tokens
+      t.integer :output_tokens
+      t.integer :cached_tokens
+      t.integer :cache_creation_tokens
+      t.timestamps
+    end
+
+    add_index :<%= message_table_name %>, :role
+  end
+end

data/lib/generators/ruby_llm/install/templates/create_models_migration.rb.tt

@@ -0,0 +1,45 @@
+class Create<%= model_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :<%= model_table_name %> do |t|
+      t.string :model_id, null: false
+      t.string :name, null: false
+      t.string :provider, null: false
+      t.string :family
+      t.datetime :model_created_at
+      t.integer :context_window
+      t.integer :max_output_tokens
+      t.date :knowledge_cutoff
+<% if postgresql? %>
+      t.jsonb :modalities, default: {}
+      t.jsonb :capabilities, default: []
+      t.jsonb :pricing, default: {}
+      t.jsonb :metadata, default: {}
+<% elsif mysql? %>
+      t.json :modalities
+      t.json :capabilities
+      t.json :pricing
+      t.json :metadata
+<% else %>
+      t.json :modalities, default: {}
+      t.json :capabilities, default: []
+      t.json :pricing, default: {}
+      t.json :metadata, default: {}
+<% end %>
+      t.timestamps
+
+      t.index [:provider, :model_id], unique: true
+      t.index :provider
+      t.index :family
+<% if postgresql? %>
+      t.index :capabilities, using: :gin
+      t.index :modalities, using: :gin
+<% end %>
+    end
+
+    # Load models from JSON
+    say_with_time "Loading models from models.json" do
+      RubyLLM.models.load_from_json!
+      <%= model_model_name %>.save_to_database
+    end
+  end
+end

data/lib/generators/ruby_llm/install/templates/create_tool_calls_migration.rb.tt

@@ -0,0 +1,20 @@
+<%#- # Migration for creating tool_calls table with database-specific JSON handling -%>
+class Create<%= tool_call_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    create_table :<%= tool_call_table_name %> do |t|
+      t.string :tool_call_id, null: false
+      t.string :name, null: false
+<% if postgresql? %>
+      t.jsonb :arguments, default: {}
+<% elsif mysql? %>
+      t.json :arguments
+<% else %>
+      t.json :arguments, default: {}
+<% end %>
+      t.timestamps
+    end
+
+    add_index :<%= tool_call_table_name %>, :tool_call_id, unique: true
+    add_index :<%= tool_call_table_name %>, :name
+  end
+end

data/lib/generators/ruby_llm/install/templates/initializer.rb.tt

@@ -0,0 +1,12 @@
+RubyLLM.configure do |config|
+  config.openai_api_key = ENV['OPENAI_API_KEY'] || Rails.application.credentials.dig(:openai_api_key)
+  # config.default_model = "gpt-4.1-nano"
+
+  # Use the new association-based acts_as API (recommended)
+  config.use_new_acts_as = true
+<% if model_model_name != 'Model' -%>
+
+  # Custom model registry class name
+  config.model_registry_class = "<%= model_model_name %>"
+<% end -%>
+end

data/lib/generators/ruby_llm/install/templates/message_model.rb.tt

@@ -0,0 +1,4 @@
+class <%= message_model_name %> < ApplicationRecord
+  <%= acts_as_message_declaration %><% unless options[:skip_active_storage] %>
+  has_many_attached :attachments<% end %>
+end

data/lib/generators/ruby_llm/install/templates/model_model.rb.tt

@@ -0,0 +1,3 @@
+class <%= model_model_name %> < ApplicationRecord
+  <%= acts_as_model_declaration %>
+end

data/lib/generators/ruby_llm/install/templates/tool_call_model.rb.tt

@@ -0,0 +1,3 @@
+class <%= tool_call_model_name %> < ApplicationRecord
+  <%= acts_as_tool_call_declaration %>
+end

data/lib/generators/ruby_llm/upgrade_to_v1_7/templates/migration.rb.tt

@@ -0,0 +1,145 @@
+class MigrateToRubyLLMModelReferences < ActiveRecord::Migration<%= migration_version %>
+  def up
+    model_class = <%= model_model_name %>
+    chat_class = <%= chat_model_name %>
+    message_class = <%= message_model_name %>
+<% if @model_table_already_existed %>
+    # Load models from models.json if Model table already existed
+    say_with_time "Loading models from models.json" do
+      RubyLLM.models.load_from_json!
+      model_class.save_to_database
+      "Loaded #{model_class.count} models"
+    end
+<% end %>
+
+    # Then check for any models in existing data that aren't in models.json
+    say_with_time "Checking for additional models in existing data" do
+      collect_and_create_models(chat_class, :<%= chat_table_name %>, model_class)
+      collect_and_create_models(message_class, :<%= message_table_name %>, model_class)
+      model_class.count
+    end
+
+    # Migrate foreign keys
+    migrate_foreign_key(:<%= chat_table_name %>, chat_class, model_class, :<%= model_table_name.singularize %>)
+    migrate_foreign_key(:<%= message_table_name %>, message_class, model_class, :<%= model_table_name.singularize %>)
+  end
+
+  def down
+    # Remove foreign key references
+    if column_exists?(:<%= message_table_name %>, :<%= model_table_name.singularize %>_id)
+      remove_reference :<%= message_table_name %>, :<%= model_table_name.singularize %>, foreign_key: true
+    end
+
+    if column_exists?(:<%= chat_table_name %>, :<%= model_table_name.singularize %>_id)
+      remove_reference :<%= chat_table_name %>, :<%= model_table_name.singularize %>, foreign_key: true
+    end
+
+    # Restore original model_id string columns
+    if column_exists?(:<%= message_table_name %>, :model_id_string)
+      rename_column :<%= message_table_name %>, :model_id_string, :model_id
+    end
+
+    if column_exists?(:<%= chat_table_name %>, :model_id_string)
+      rename_column :<%= chat_table_name %>, :model_id_string, :model_id
+    end
+  end
+
+  private
+
+  def collect_and_create_models(record_class, table_name, model_class)
+    return unless column_exists?(table_name, :model_id)
+
+    has_provider = column_exists?(table_name, :provider)
+
+    # Collect unique model/provider combinations using read_attribute to bypass overrides
+    models_set = Set.new
+
+    record_class.find_each do |record|
+      model_id = record.read_attribute(:model_id)
+      next if model_id.blank?
+
+      provider = has_provider ? record.read_attribute(:provider) : nil
+      models_set.add([ model_id, provider ])
+    end
+
+    models_set.each do |model_id, provider|
+      find_or_create_model(model_id, provider, model_class)
+    end
+  end
+
+  def find_or_create_model(model_id, provider, model_class)
+    return if model_id.blank?
+
+    begin
+      model_info, _provider = RubyLLM.models.resolve(model_id, provider: provider)
+
+      model_class.find_or_create_by!(
+        model_id: model_info.id,
+        provider: model_info.provider
+      ) do |m|
+        m.name = model_info.name || model_info.id
+        m.family = model_info.family
+        m.model_created_at = model_info.created_at
+        m.context_window = model_info.context_window
+        m.max_output_tokens = model_info.max_output_tokens
+        m.knowledge_cutoff = model_info.knowledge_cutoff
+        m.modalities = model_info.modalities.to_h
+        m.capabilities = model_info.capabilities
+        m.pricing = model_info.pricing.to_h
+        m.metadata = model_info.metadata
+      end
+    rescue => e
+      # Skip models that can't be resolved - they'll need manual fixing
+      Rails.logger.warn "Skipping unresolvable model: #{model_id} - will need manual update"
+      nil
+    end
+  end
+
+
+  def migrate_foreign_key(table_name, record_class, model_class, foreign_key_name)
+    return unless column_exists?(table_name, :model_id)
+
+    # Check if we need to rename the string column to avoid collision
+    if column_exists?(table_name, :model_id) && !foreign_key_exists?(table_name, :models)
+      # Temporarily rename the string column
+      rename_column table_name, :model_id, :model_id_string
+    end
+
+    # Add the foreign key reference
+    unless column_exists?(table_name, "#{foreign_key_name}_id")
+      add_reference table_name, foreign_key_name, foreign_key: true
+    end
+
+    say_with_time "Migrating #{table_name} model references" do
+      record_class.reset_column_information
+      has_provider = column_exists?(table_name, :provider)
+
+      # Determine which column to read from (renamed or original)
+      model_id_column = column_exists?(table_name, :model_id_string) ? :model_id_string : :model_id
+
+      record_class.find_each do |record|
+        model_id = record.read_attribute(model_id_column)
+        next if model_id.blank?
+
+        provider = has_provider ? record.read_attribute(:provider) : nil
+
+        model = if has_provider && provider.present?
+          model_class.find_by(model_id: model_id, provider: provider)
+        else
+          find_model_for_record(model_id, model_class)
+        end
+
+        record.update_column("#{foreign_key_name}_id", model.id) if model
+      end
+    end
+  end
+
+  def find_model_for_record(model_id, model_class)
+    begin
+      model_info, _provider = RubyLLM.models.resolve(model_id)
+      model_class.find_by(model_id: model_info.id, provider: model_info.provider)
+    rescue => e
+      model_class.find_by(model_id: model_id)
+    end
+  end
+end