dify_llm 1.6.4 → 1.7.1

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

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+
+require 'rails/generators'
+require 'rails/generators/active_record'
+require_relative '../generator_helpers'
+
+module RubyLLM
+  # Generator for RubyLLM Rails models and migrations
+  class InstallGenerator < Rails::Generators::Base
+    include Rails::Generators::Migration
+    include RubyLLM::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
+      # Create migrations with timestamps to ensure proper order
+      # First create chats table
+      migration_template 'create_chats_migration.rb.tt',
+                         "db/migrate/create_#{chat_table_name}.rb"
+
+      # Then create messages table (must come before tool_calls due to foreign key)
+      sleep 1 # Ensure different timestamp
+      migration_template 'create_messages_migration.rb.tt',
+                         "db/migrate/create_#{message_table_name}.rb"
+
+      # Then create tool_calls table (references messages)
+      sleep 1 # Ensure different timestamp
+      migration_template 'create_tool_calls_migration.rb.tt',
+                         "db/migrate/create_#{tool_call_table_name}.rb"
+
+      # Create models table
+      sleep 1 # Ensure different timestamp
+      migration_template 'create_models_migration.rb.tt',
+                         "db/migrate/create_#{model_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

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

@@ -1,3 +1,3 @@
-class <%= options[:chat_model_name] %> < ApplicationRecord
+class <%= chat_model_name %> < ApplicationRecord
   <%= acts_as_chat_declaration %>
 end

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

@@ -1,7 +1,7 @@
-class Create<%= options[:chat_model_name].pluralize %> < ActiveRecord::Migration<%= migration_version %>
+class Create<%= chat_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
   def change
-    create_table :<%= options[:chat_model_name].tableize %> do |t|
-      t.references :<%= options[:model_model_name].tableize.singularize %>, foreign_key: true
+    create_table :<%= chat_table_name %> do |t|
+      t.references :<%= model_table_name.singularize %>, foreign_key: true
       t.timestamps
     end
   end

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

@@ -1,16 +1,16 @@
-class Create<%= options[:message_model_name].pluralize %> < ActiveRecord::Migration<%= migration_version %>
+class Create<%= message_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
   def change
-    create_table :<%= options[:message_model_name].tableize %> do |t|
-      t.references :<%= options[:chat_model_name].tableize.singularize %>, null: false, foreign_key: true
+    create_table :<%= message_table_name %> do |t|
+      t.references :<%= chat_table_name.singularize %>, null: false, foreign_key: true
       t.string :role, null: false
       t.text :content
-      t.references :<%= options[:model_model_name].tableize.singularize %>, foreign_key: true
+      t.references :<%= model_table_name.singularize %>, foreign_key: true
       t.integer :input_tokens
       t.integer :output_tokens
-      t.references :<%= options[:tool_call_model_name].tableize.singularize %>, foreign_key: true
+      t.references :<%= tool_call_table_name.singularize %>, foreign_key: true
       t.timestamps
     end
 
-    add_index :<%= options[:message_model_name].tableize %>, :role
+    add_index :<%= message_table_name %>, :role
   end
 end

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

@@ -1,6 +1,6 @@
-class Create<%= options[:model_model_name].pluralize %> < ActiveRecord::Migration<%= migration_version %>
+class Create<%= model_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
   def change
-    create_table :<%= options[:model_model_name].tableize %> do |t|
+    create_table :<%= model_table_name %> do |t|
       t.string :model_id, null: false
       t.string :name, null: false
       t.string :provider, null: false
@@ -34,10 +34,7 @@ class Create<%= options[:model_model_name].pluralize %> < ActiveRecord::Migratio
     # Load models from JSON
     say_with_time "Loading models from models.json" do
       RubyLLM.models.load_from_json!
-      model_class = '<%= options[:model_model_name] %>'.constantize
-      model_class.save_to_database
-
-      "Loaded #{model_class.count} models"
+      <%= model_model_name %>.save_to_database
     end
   end
 end

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

@@ -1,15 +1,15 @@
 <%#- # Migration for creating tool_calls table with database-specific JSON handling -%>
-class Create<%= options[:tool_call_model_name].pluralize %> < ActiveRecord::Migration<%= migration_version %>
+class Create<%= tool_call_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
   def change
-    create_table :<%= options[:tool_call_model_name].tableize %> do |t|
-      t.references :<%= options[:message_model_name].tableize.singularize %>, null: false, foreign_key: true
+    create_table :<%= tool_call_table_name %> do |t|
+      t.references :<%= message_table_name.singularize %>, null: false, foreign_key: true
       t.string :tool_call_id, null: false
       t.string :name, null: false
       t.<%= postgresql? ? 'jsonb' : 'json' %> :arguments, default: {}
       t.timestamps
     end
 
-    add_index :<%= options[:tool_call_model_name].tableize %>, :tool_call_id, unique: true
-    add_index :<%= options[:tool_call_model_name].tableize %>, :name
+    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

@@ -2,8 +2,11 @@ RubyLLM.configure do |config|
   config.openai_api_key = Rails.application.credentials.dig(:openai_api_key)
   # config.default_model = "gpt-4.1-nano"
 
-<% unless skip_model_registry? -%>
-  # DB-backed model registry for rich model metadata
-  config.model_registry_class = "<%= options[:model_model_name] %>"
+  # 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

@@ -1,4 +1,4 @@
-class <%= options[:message_model_name] %> < ApplicationRecord
+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

@@ -1,3 +1,3 @@
-class <%= options[:model_model_name] %> < ApplicationRecord
+class <%= model_model_name %> < ApplicationRecord
   <%= acts_as_model_declaration %>
 end

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

@@ -1,3 +1,3 @@
-class <%= options[:tool_call_model_name] %> < ApplicationRecord
+class <%= tool_call_model_name %> < ApplicationRecord
   <%= acts_as_tool_call_declaration %>
 end

data/lib/generators/ruby_llm/{migrate_model_fields → upgrade_to_v1_7}/templates/migration.rb.tt

@@ -1,43 +1,46 @@
 class MigrateToRubyLLMModelReferences < ActiveRecord::Migration<%= migration_version %>
   def up
-    # Ensure Models table exists
-    unless table_exists?(:<%= options[:model_model_name].tableize %>)
-      raise "Models table doesn't exist. Please run 'rails generate ruby_llm:install' with model registry first."
+    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
-
-    model_class = <%= options[:model_model_name] %>
-    chat_class = <%= options[:chat_model_name] %>
-    message_class = <%= options[:message_model_name] %>
+<% 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, :<%= options[:chat_model_name].tableize %>, model_class)
-      collect_and_create_models(message_class, :<%= options[:message_model_name].tableize %>, model_class)
+      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(:<%= options[:chat_model_name].tableize %>, chat_class, model_class, :<%= options[:model_model_name].underscore %>)
-    migrate_foreign_key(:<%= options[:message_model_name].tableize %>, message_class, model_class, :<%= options[:model_model_name].underscore %>)
+    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?(:<%= options[:message_model_name].tableize %>, :<%= options[:model_model_name].underscore %>_id)
-      remove_reference :<%= options[:message_model_name].tableize %>, :<%= options[:model_model_name].underscore %>, foreign_key: true
+    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?(:<%= options[:chat_model_name].tableize %>, :<%= options[:model_model_name].underscore %>_id)
-      remove_reference :<%= options[:chat_model_name].tableize %>, :<%= options[:model_model_name].underscore %>, foreign_key: true
+    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?(:<%= options[:message_model_name].tableize %>, :model_id_string)
-      rename_column :<%= options[:message_model_name].tableize %>, :model_id_string, :model_id
+    if column_exists?(:<%= message_table_name %>, :model_id_string)
+      rename_column :<%= message_table_name %>, :model_id_string, :model_id
     end
 
-    if column_exists?(:<%= options[:chat_model_name].tableize %>, :model_id_string)
-      rename_column :<%= options[:chat_model_name].tableize %>, :model_id_string, :model_id
+    if column_exists?(:<%= chat_table_name %>, :model_id_string)
+      rename_column :<%= chat_table_name %>, :model_id_string, :model_id
     end
   end
 
@@ -139,4 +142,4 @@ class MigrateToRubyLLMModelReferences < ActiveRecord::Migration<%= migration_ver
       model_class.find_by(model_id: model_id)
     end
   end
-end
+end

data/lib/generators/ruby_llm/upgrade_to_v1_7/upgrade_to_v1_7_generator.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+require 'rails/generators'
+require 'rails/generators/active_record'
+require_relative '../generator_helpers'
+
+module RubyLLM
+  class UpgradeToV17Generator < Rails::Generators::Base # rubocop:disable Style/Documentation
+    include Rails::Generators::Migration
+    include RubyLLM::GeneratorHelpers
+
+    namespace 'ruby_llm:upgrade_to_v1_7'
+    source_root File.expand_path('templates', __dir__)
+
+    # Override source_paths to include install templates
+    def self.source_paths
+      [
+        File.expand_path('templates', __dir__),
+        File.expand_path('../install/templates', __dir__)
+      ]
+    end
+
+    argument :model_mappings, type: :array, default: [], banner: 'chat:ChatName message:MessageName ...'
+
+    desc 'Upgrades existing RubyLLM apps to v1.7 with new Rails-like API\n' \
+         'Usage: rails g ruby_llm:upgrade_to_v1_7 [chat:ChatName] [message:MessageName] ...'
+
+    def self.next_migration_number(dirname)
+      ::ActiveRecord::Generators::Base.next_migration_number(dirname)
+    end
+
+    def create_migration_file
+      @model_table_already_existed = table_exists?(table_name_for(model_model_name))
+
+      # First check if models table exists, if not create it
+      unless @model_table_already_existed
+        migration_template 'create_models_migration.rb.tt',
+                           "db/migrate/create_#{table_name_for(model_model_name)}.rb",
+                           migration_version: migration_version,
+                           model_model_name: model_model_name
+
+        sleep 1 # Ensure different timestamp
+      end
+
+      migration_template 'migration.rb.tt',
+                         'db/migrate/migrate_to_ruby_llm_model_references.rb',
+                         migration_version: migration_version,
+                         chat_model_name: chat_model_name,
+                         message_model_name: message_model_name,
+                         tool_call_model_name: tool_call_model_name,
+                         model_model_name: model_model_name,
+                         model_table_already_existed: @model_table_already_existed
+    end
+
+    def create_model_file
+      create_namespace_modules
+
+      template 'model_model.rb.tt', "app/models/#{model_model_name.underscore}.rb"
+    end
+
+    def update_existing_models
+      update_model_acts_as(chat_model_name, 'acts_as_chat', acts_as_chat_declaration)
+      update_model_acts_as(message_model_name, 'acts_as_message', acts_as_message_declaration)
+      update_model_acts_as(tool_call_model_name, 'acts_as_tool_call', acts_as_tool_call_declaration)
+    end
+
+    def update_initializer
+      initializer_path = 'config/initializers/ruby_llm.rb'
+
+      unless File.exist?(initializer_path)
+        say_status :warning, 'No initializer found. Creating one...', :yellow
+        template 'initializer.rb.tt', initializer_path
+        return
+      end
+
+      initializer_content = File.read(initializer_path)
+
+      return if initializer_content.include?('config.use_new_acts_as')
+
+      inject_into_file initializer_path, before: /^end/ do
+        lines = ["\n  # Enable the new Rails-like API", '  config.use_new_acts_as = true']
+        lines << "  config.model_registry_class = \"#{model_model_name}\"" if model_model_name != 'Model'
+        lines << "\n"
+        lines.join("\n")
+      end
+    end
+
+    def show_next_steps
+      say_status :success, 'Upgrade prepared!', :green
+      say <<~INSTRUCTIONS
+
+        Next steps:
+        1. Review the generated migrations
+        2. Run: rails db:migrate
+        3. Update your code to use the new API: #{chat_model_name}.create! now has the same signature as RubyLLM.chat
+
+        ⚠️  If you get "undefined method 'acts_as_model'" during migration:
+           Add this to config/application.rb BEFORE your Application class:
+
+           RubyLLM.configure do |config|
+             config.use_new_acts_as = true
+           end
+
+        📚 See the full migration guide: https://rubyllm.com/upgrading-to-1-7/
+
+      INSTRUCTIONS
+    end
+
+    private
+
+    def update_model_acts_as(model_name, old_acts_as, new_acts_as)
+      model_path = "app/models/#{model_name.underscore}.rb"
+      return unless File.exist?(Rails.root.join(model_path))
+
+      content = File.read(Rails.root.join(model_path))
+      return unless content.match?(/^\s*#{old_acts_as}/)
+
+      gsub_file model_path, /^\s*#{old_acts_as}.*$/, "  #{new_acts_as}"
+    end
+  end
+end

data/lib/ruby_llm/active_record/acts_as.rb

@@ -31,105 +31,135 @@ module RubyLLM
       end
 
       class_methods do # rubocop:disable Metrics/BlockLength
-        def acts_as_chat(message_class: 'Message', tool_call_class: 'ToolCall',
-                         model_class: 'Model', model_foreign_key: nil)
+        def acts_as_chat(messages: :messages, message_class: nil,
+                         model: :model, model_class: nil)
           include RubyLLM::ActiveRecord::ChatMethods
 
-          @message_class = message_class.to_s
-          @tool_call_class = tool_call_class.to_s
-          @model_class = model_class.to_s
-          @model_foreign_key = model_foreign_key || ActiveSupport::Inflector.foreign_key(@model_class)
+          class_attribute :messages_association_name, :model_association_name, :message_class, :model_class
 
-          has_many :messages,
+          self.messages_association_name = messages
+          self.model_association_name = model
+          self.message_class = (message_class || messages.to_s.classify).to_s
+          self.model_class = (model_class || model.to_s.classify).to_s
+
+          has_many messages,
                    -> { order(created_at: :asc) },
-                   class_name: @message_class,
-                   inverse_of: :chat,
+                   class_name: self.message_class,
                    dependent: :destroy
 
-          belongs_to :model,
-                     class_name: @model_class,
-                     foreign_key: @model_foreign_key,
+          belongs_to model,
+                     class_name: self.model_class,
                      optional: true
 
           delegate :add_message, to: :to_llm
+
+          define_method :messages_association do
+            send(messages_association_name)
+          end
+
+          define_method :model_association do
+            send(model_association_name)
+          end
+
+          define_method :'model_association=' do |value|
+            send("#{model_association_name}=", value)
+          end
         end
 
-        def acts_as_model(chat_class: 'Chat')
+        def acts_as_model(chats: :chats, chat_class: nil)
           include RubyLLM::ActiveRecord::ModelMethods
 
-          @chat_class = chat_class.to_s
+          class_attribute :chats_association_name, :chat_class
+
+          self.chats_association_name = chats
+          self.chat_class = (chat_class || chats.to_s.classify).to_s
 
           validates :model_id, presence: true, uniqueness: { scope: :provider }
           validates :provider, presence: true
           validates :name, presence: true
 
-          has_many :chats,
-                   class_name: @chat_class,
-                   foreign_key: ActiveSupport::Inflector.foreign_key(name)
+          has_many chats, class_name: self.chat_class
+
+          define_method :chats_association do
+            send(chats_association_name)
+          end
         end
 
-        def acts_as_message(chat_class: 'Chat', # rubocop:disable Metrics/ParameterLists
-                            chat_foreign_key: nil,
-                            tool_call_class: 'ToolCall',
-                            tool_call_foreign_key: nil,
-                            model_class: 'Model',
-                            model_foreign_key: nil,
-                            touch_chat: false)
+        def acts_as_message(chat: :chat, chat_class: nil, touch_chat: false, # rubocop:disable Metrics/ParameterLists
+                            tool_calls: :tool_calls, tool_call_class: nil,
+                            model: :model, model_class: nil)
           include RubyLLM::ActiveRecord::MessageMethods
 
-          @chat_class = chat_class.to_s
-          @chat_foreign_key = chat_foreign_key || ActiveSupport::Inflector.foreign_key(@chat_class)
-
-          @tool_call_class = tool_call_class.to_s
-          @tool_call_foreign_key = tool_call_foreign_key || ActiveSupport::Inflector.foreign_key(@tool_call_class)
+          class_attribute :chat_association_name, :tool_calls_association_name, :model_association_name,
+                          :chat_class, :tool_call_class, :model_class
 
-          @model_class = model_class.to_s
-          @model_foreign_key = model_foreign_key || ActiveSupport::Inflector.foreign_key(@model_class)
+          self.chat_association_name = chat
+          self.tool_calls_association_name = tool_calls
+          self.model_association_name = model
+          self.chat_class = (chat_class || chat.to_s.classify).to_s
+          self.tool_call_class = (tool_call_class || tool_calls.to_s.classify).to_s
+          self.model_class = (model_class || model.to_s.classify).to_s
 
-          belongs_to :chat,
-                     class_name: @chat_class,
-                     foreign_key: @chat_foreign_key,
-                     inverse_of: :messages,
+          belongs_to chat,
+                     class_name: self.chat_class,
                      touch: touch_chat
 
-          has_many :tool_calls,
-                   class_name: @tool_call_class,
+          has_many tool_calls,
+                   class_name: self.tool_call_class,
                    dependent: :destroy
 
           belongs_to :parent_tool_call,
-                     class_name: @tool_call_class,
-                     foreign_key: @tool_call_foreign_key,
-                     optional: true,
-                     inverse_of: :result
+                     class_name: self.tool_call_class,
+                     foreign_key: ActiveSupport::Inflector.foreign_key(tool_calls.to_s.singularize),
+                     optional: true
 
           has_many :tool_results,
-                   through: :tool_calls,
+                   through: tool_calls,
                    source: :result,
-                   class_name: @message_class
+                   class_name: name
 
-          belongs_to :model,
-                     class_name: @model_class,
-                     foreign_key: @model_foreign_key,
+          belongs_to model,
+                     class_name: self.model_class,
                      optional: true
 
           delegate :tool_call?, :tool_result?, to: :to_llm
+
+          define_method :chat_association do
+            send(chat_association_name)
+          end
+
+          define_method :tool_calls_association do
+            send(tool_calls_association_name)
+          end
+
+          define_method :model_association do
+            send(model_association_name)
+          end
         end
 
-        def acts_as_tool_call(message_class: 'Message', message_foreign_key: nil, result_foreign_key: nil)
-          @message_class = message_class.to_s
-          @message_foreign_key = message_foreign_key || ActiveSupport::Inflector.foreign_key(@message_class)
-          @result_foreign_key = result_foreign_key || ActiveSupport::Inflector.foreign_key(name)
+        def acts_as_tool_call(message: :message, message_class: nil,
+                              result: :result, result_class: nil)
+          class_attribute :message_association_name, :result_association_name, :message_class, :result_class
 
-          belongs_to :message,
-                     class_name: @message_class,
-                     foreign_key: @message_foreign_key,
-                     inverse_of: :tool_calls
+          self.message_association_name = message
+          self.result_association_name = result
+          self.message_class = (message_class || message.to_s.classify).to_s
+          self.result_class = (result_class || self.message_class).to_s
 
-          has_one :result,
-                  class_name: @message_class,
-                  foreign_key: @result_foreign_key,
-                  inverse_of: :parent_tool_call,
+          belongs_to message,
+                     class_name: self.message_class
+
+          has_one result,
+                  class_name: self.result_class,
                   dependent: :nullify
+
+          define_method :message_association do
+            send(message_association_name)
+          end
+
+          define_method :result_association do
+            send(result_association_name)
+          end
         end
       end
     end