diff --git a/lib/ruby_llm/active_record/acts_as.rb b/lib/ruby_llm/active_record/acts_as.rb index 9a9ecbf5f..408fb2640 100644 --- a/lib/ruby_llm/active_record/acts_as.rb +++ b/lib/ruby_llm/active_record/acts_as.rb @@ -28,28 +28,41 @@ def acts_as_message(chat_class: 'Chat', tool_call_class: 'ToolCall', touch_chat: include MessageMethods @chat_class = chat_class.to_s + @chat_foreign_key = "#{@chat_class.underscore}_id" @tool_call_class = tool_call_class.to_s + @tool_call_foreign_key = "#{@tool_call_class.underscore}_id" - belongs_to :chat, class_name: @chat_class, touch: touch_chat - has_many :tool_calls, class_name: @tool_call_class, dependent: :destroy + belongs_to :chat, + class_name: @chat_class, + foreign_key: @chat_foreign_key, + inverse_of: :messages, + touch: touch_chat + + has_many :tool_calls, + class_name: @tool_call_class, + dependent: :destroy belongs_to :parent_tool_call, class_name: @tool_call_class, - foreign_key: 'tool_call_id', + foreign_key: @tool_call_foreign_key, optional: true, inverse_of: :result delegate :tool_call?, :tool_result?, :tool_results, to: :to_llm end - def acts_as_tool_call(message_class: 'Message') + def acts_as_tool_call(message_class: 'Message') # rubocop:disable Metrics/MethodLength @message_class = message_class.to_s + @message_foreign_key = "#{@message_class.underscore}_id" - belongs_to :message, class_name: @message_class + belongs_to :message, + class_name: @message_class, + foreign_key: @message_foreign_key, + inverse_of: :tool_calls has_one :result, class_name: @message_class, - foreign_key: 'tool_call_id', + foreign_key: @message_foreign_key, inverse_of: :parent_tool_call, dependent: :nullify end @@ -159,14 +172,15 @@ def persist_message_completion(message) # rubocop:disable Metrics/AbcSize,Metric end transaction do - @message.update!( + @message.update( role: message.role, content: message.content, model_id: message.model_id, - tool_call_id: tool_call_id, input_tokens: message.input_tokens, output_tokens: message.output_tokens ) + @message.write_attribute(@message.class.tool_call_foreign_key, tool_call_id) if tool_call_id + @message.save! persist_tool_calls(message.tool_calls) if message.tool_calls.present? end end @@ -185,6 +199,11 @@ def persist_tool_calls(tool_calls) module MessageMethods extend ActiveSupport::Concern + class_methods do + attr_reader :chat_class, :tool_call_class + attr_reader :chat_foreign_key, :tool_call_foreign_key + end + def to_llm RubyLLM::Message.new( role: role.to_sym, diff --git a/spec/ruby_llm/active_record/acts_as_spec.rb b/spec/ruby_llm/active_record/acts_as_spec.rb index b8d8bf6c3..9715d5352 100644 --- a/spec/ruby_llm/active_record/acts_as_spec.rb +++ b/spec/ruby_llm/active_record/acts_as_spec.rb @@ -19,6 +19,16 @@ t.timestamps end + # the Bot* classes are used to test the class + # renaming functionality of acts_as_* + # They are supposed to be identical to the + # non-Bot* classes, but with a different names + # using Rails-canonical naming conventions. + create_table :bot_chats do |t| + t.string :model_id + t.timestamps + end + create_table :messages do |t| t.references :chat t.string :role @@ -30,6 +40,17 @@ t.timestamps end + create_table :bot_messages do |t| + t.references :bot_chat + t.string :role + t.text :content + t.string :model_id + t.integer :input_tokens + t.integer :output_tokens + t.references :bot_tool_call + t.timestamps + end + create_table :tool_calls do |t| t.references :message t.string :tool_call_id @@ -37,6 +58,14 @@ t.json :arguments t.timestamps end + + create_table :bot_tool_calls do |t| + t.references :bot_message + t.string :tool_call_id + t.string :name + t.json :arguments + t.timestamps + end end end @@ -45,16 +74,31 @@ class Chat < ActiveRecord::Base # rubocop:disable Lint/ConstantDefinitionInBlock acts_as_chat end + class BotChat < ActiveRecord::Base # rubocop:disable Lint/ConstantDefinitionInBlock,RSpec/LeakyConstantDeclaration + include RubyLLM::ActiveRecord::ActsAs + acts_as_chat message_class: 'BotMessage', tool_call_class: 'BotToolCall' + end + class Message < ActiveRecord::Base # rubocop:disable Lint/ConstantDefinitionInBlock,RSpec/LeakyConstantDeclaration include RubyLLM::ActiveRecord::ActsAs acts_as_message end + class BotMessage < ActiveRecord::Base # rubocop:disable Lint/ConstantDefinitionInBlock,RSpec/LeakyConstantDeclaration + include RubyLLM::ActiveRecord::ActsAs + acts_as_message chat_class: 'BotChat', tool_call_class: 'BotToolCall' + end + class ToolCall < ActiveRecord::Base # rubocop:disable Lint/ConstantDefinitionInBlock,RSpec/LeakyConstantDeclaration include RubyLLM::ActiveRecord::ActsAs acts_as_tool_call end + class BotToolCall < ActiveRecord::Base # rubocop:disable Lint/ConstantDefinitionInBlock,RSpec/LeakyConstantDeclaration + include RubyLLM::ActiveRecord::ActsAs + acts_as_tool_call message_class: 'BotMessage' + end + class Calculator < RubyLLM::Tool # rubocop:disable Lint/ConstantDefinitionInBlock,RSpec/LeakyConstantDeclaration description 'Performs basic arithmetic' @@ -71,19 +115,23 @@ def execute(expression:) shared_examples 'a chainable chat method' do |method_name, *args| it "returns a Chat instance for ##{method_name}" do - chat = Chat.create!(model_id: 'gpt-4.1-nano') - result = chat.public_send(method_name, *args) - expect(result).to be_a(Chat) + [Chat, BotChat].each do |chat_class| + chat = chat_class.create!(model_id: 'gpt-4.1-nano') + result = chat.public_send(method_name, *args) + expect(result).to be_a(chat_class) + end end end shared_examples 'a chainable callback method' do |callback_name| - it "supports #{callback_name} callback" do - chat = Chat.create!(model_id: 'gpt-4.1-nano') - result = chat.public_send(callback_name) do - # no-op for testing + it "supports #{callback_name} callback" do # rubocop:disable RSpec/ExampleLength + [Chat, BotChat].each do |chat_class| + chat = chat_class.create!(model_id: 'gpt-4.1-nano') + result = chat.public_send(callback_name) do + # no-op for testing + end + expect(result).to be_a(chat_class) end - expect(result).to be_a(Chat) end end @@ -124,9 +172,11 @@ def execute(expression:) describe 'with_tools functionality' do it 'returns a Chat instance when using with_tool' do - chat = Chat.create!(model_id: 'gpt-4.1-nano') - with_tool_result = chat.with_tool(Calculator) - expect(with_tool_result).to be_a(Chat) + [Chat, BotChat].each do |chat_class| + chat = chat_class.create!(model_id: 'gpt-4.1-nano') + with_tool_result = chat.with_tool(Calculator) + expect(with_tool_result).to be_a(chat_class) + end end it 'persists user messages' do @@ -145,11 +195,13 @@ def execute(expression:) it_behaves_like 'a chainable callback method', :on_new_message it_behaves_like 'a chainable callback method', :on_end_message - it 'supports method chaining with tools' do - chat = Chat.create!(model_id: 'gpt-4.1-nano') - chat.with_tool(Calculator) - .with_temperature(0.5) - expect(chat).to be_a(Chat) + it 'supports method chaining with tools' do # rubocop:disable RSpec/ExampleLength + [Chat, BotChat].each do |chat_class| + chat = chat_class.create!(model_id: 'gpt-4.1-nano') + chat.with_tool(Calculator) + .with_temperature(0.5) + expect(chat).to be_a(chat_class) + end end it 'persists messages after chaining' do