Merge branch 'main' into fix-models-resolve-delegation

Author: crmne

lib/generators/ruby_llm/install/install_generator.rb

@@ -32,12 +32,12 @@ def create_migration_files
       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)
+      # Then create messages table
       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)
+      # Then create tool_calls table
       sleep 1 # Ensure different timestamp
       migration_template 'create_tool_calls_migration.rb.tt',
                          "db/migrate/create_#{tool_call_table_name}.rb"
@@ -46,6 +46,12 @@ def create_migration_files
       sleep 1 # Ensure different timestamp
       migration_template 'create_models_migration.rb.tt',
                          "db/migrate/create_#{model_table_name}.rb"
+
+      # Add references to chats, tool_calls and messages.
+      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

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

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

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

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

@@ -1,13 +1,10 @@
 class Create<%= message_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
   def change
     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 :<%= model_table_name.singularize %>, foreign_key: true
       t.integer :input_tokens
       t.integer :output_tokens
-      t.references :<%= tool_call_table_name.singularize %>, foreign_key: true
       t.timestamps
     end
 

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

@@ -2,7 +2,6 @@
 class Create<%= tool_call_model_name.gsub('::', '').pluralize %> < ActiveRecord::Migration<%= migration_version %>
   def change
     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: {}

spec/lib/generators/ruby_llm/install_generator_spec.rb

@@ -15,7 +15,8 @@
         'create_chats_migration.rb.tt',
         'create_messages_migration.rb.tt',
         'create_tool_calls_migration.rb.tt',
-        'create_models_migration.rb.tt'
+        'create_models_migration.rb.tt',
+        'add_references_to_chats_tool_calls_and_messages_migration.rb.tt'
       ]
     end
 
@@ -31,10 +32,6 @@
       it 'defines chats table' do
         expect(chat_migration).to include('create_table :<%= chat_table_name %>')
       end
-
-      it 'includes model reference' do
-        expect(chat_migration).to include('t.references :<%= model_table_name.singularize %>')
-      end
     end
 
     describe 'messages migration' do
@@ -44,10 +41,6 @@
         expect(message_migration).to include('create_table :<%= message_table_name %>')
       end
 
-      it 'includes chat reference' do
-        expect(message_migration).to include('t.references :<%= chat_table_name.singularize %>, null: false, foreign_key: true') # rubocop:disable Layout/LineLength
-      end
-
       it 'includes role field' do
         expect(message_migration).to include('t.string :role')
       end
@@ -72,6 +65,39 @@
         expect(tool_call_migration).to include('t.string :name')
       end
     end
+
+    describe 'add references migration' do
+      let(:add_references_migration) do
+        File.read(File.join(template_dir, 'add_references_to_chats_tool_calls_and_messages_migration.rb.tt'))
+      end
+
+      it 'adds model reference to chats' do
+        expect(add_references_migration).to include('add_reference :<%= chat_table_name %>, ' \
+                                                    ':<%= model_table_name.singularize %>, foreign_key: true')
+      end
+
+      it 'adds message reference to tool_calls' do
+        expect(add_references_migration).to include('add_reference :<%= tool_call_table_name %>, ' \
+                                                    ':<%= message_table_name.singularize %>, ' \
+                                                    'null: false, foreign_key: true')
+      end
+
+      it 'adds chat reference to messages' do
+        expect(add_references_migration).to include('add_reference :<%= message_table_name %>, ' \
+                                                    ':<%= chat_table_name.singularize %>, ' \
+                                                    'null: false, foreign_key: true')
+      end
+
+      it 'adds model reference to messages' do
+        expect(add_references_migration).to include('add_reference :<%= message_table_name %>, ' \
+                                                    ':<%= model_table_name.singularize %>, foreign_key: true')
+      end
+
+      it 'adds tool_call reference to messages' do
+        expect(add_references_migration).to include('add_reference :<%= message_table_name %>, ' \
+                                                    ':<%= tool_call_table_name.singularize %>, foreign_key: true')
+      end
+    end
   end
 
   describe 'JSON handling in migrations' do
@@ -257,10 +283,18 @@
       expect(models_position).to be > tool_calls_position if models_position
     end
 
-    it 'has comments explaining the order' do
+    it 'adds references after creating all tables' do
       migration_section = generator_content[/def create_migration_files.*?\n    end/m]
-      expect(migration_section).to include('must come before tool_calls due to foreign key')
-      expect(migration_section).to include('references messages')
+
+      add_references_position = migration_section.index(
+        'add_references_to_chats_tool_calls_and_messages_migration.rb.tt'
+      )
+      models_position = migration_section.index('model_table_name')
+
+      expect(add_references_position).not_to be_nil
+      expect(models_position).not_to be_nil
+
+      expect(add_references_position).to be > models_position
     end
   end