diff --git a/lib/ruby_llm/active_record/chat_methods.rb b/lib/ruby_llm/active_record/chat_methods.rb index 114d7ff7d..41d9a4447 100644 --- a/lib/ruby_llm/active_record/chat_methods.rb +++ b/lib/ruby_llm/active_record/chat_methods.rb @@ -86,7 +86,7 @@ def to_llm ) @chat.reset_messages! - ordered_messages = order_messages_for_llm(messages_association.to_a) + ordered_messages = order_messages_for_llm(eager_load_messages) ordered_messages.each do |msg| @chat.add_message(msg.to_llm) end @@ -248,7 +248,7 @@ def cleanup_failed_messages def cleanup_orphaned_tool_results # rubocop:disable Metrics/PerceivedComplexity messages_association.reload - last = messages_association.order(:id).last + last = eager_load_messages.max_by(&:id) return unless last&.tool_call? || last&.tool_result? @@ -267,6 +267,16 @@ def cleanup_orphaned_tool_results # rubocop:disable Metrics/PerceivedComplexity end end + def eager_load_messages + assoc = messages_association + return assoc.to_a unless assoc.respond_to?(:includes) + + msg_class = assoc.klass + tool_calls_name = msg_class.tool_calls_association_name + model_name = msg_class.model_association_name + assoc.includes(tool_calls_name, :parent_tool_call, model_name).to_a + end + def setup_persistence_callbacks return @chat if @chat.instance_variable_get(:@_persistence_callbacks_setup) diff --git a/spec/ruby_llm/active_record/acts_as_spec.rb b/spec/ruby_llm/active_record/acts_as_spec.rb index 6d2f3e857..3283dc6e5 100644 --- a/spec/ruby_llm/active_record/acts_as_spec.rb +++ b/spec/ruby_llm/active_record/acts_as_spec.rb @@ -1029,6 +1029,43 @@ class ToolCall < ActiveRecord::Base # rubocop:disable RSpec/LeakyConstantDeclara end end + describe 'strict_loading compatibility' do + # Verify that to_llm and cleanup_orphaned_tool_results eager-load message + # associations, preventing N+1 queries with strict_loading enabled. + + let(:chat_with_tool_calls) do + chat = Chat.create!(model: model) + chat.messages.create!(role: 'user', content: "What's 2 + 2?") + assistant_msg = chat.messages.create!(role: 'assistant', content: nil) + tool_call = assistant_msg.tool_calls.create!( + tool_call_id: 'call_strict_1', + name: 'calculator', + arguments: { expression: '2 + 2' }.to_json + ) + chat.messages.create!(role: 'tool', content: '4', parent_tool_call: tool_call) + chat.messages.create!(role: 'assistant', content: 'The answer is 4.') + chat + end + + around do |example| + chat_with_tool_calls # create data before enabling strict_loading + ApplicationRecord.strict_loading_by_default = true + ApplicationRecord.strict_loading_mode = :n_plus_one_only + example.run + ensure + ApplicationRecord.strict_loading_by_default = false + ApplicationRecord.strict_loading_mode = :all + end + + it 'to_llm does not raise StrictLoadingViolationError' do + expect { chat_with_tool_calls.reload.to_llm }.not_to raise_error + end + + it 'cleanup_orphaned_tool_results does not raise StrictLoadingViolationError' do + expect { chat_with_tool_calls.reload.send(:cleanup_orphaned_tool_results) }.not_to raise_error + end + end + describe 'extended thinking persistence' do def thinking_config_for(provider) case provider