Add API for setting messages via ActiveRecord Relations

README.md

@@ -173,7 +173,52 @@ This approach lets you recreate a conversation's history (perhaps from your data
 ## Getting and setting messages directly
 
 - You can call `.messages` to get an array containing the conversation so far.
-- TODO: Setting `.messages` will replace the conversation with the provided array.
+- You can set `.messages` to replace the conversation with a provided array or ActiveRecord::Relation:
+
+```ruby
+# Create a new chat instance
+chat = AI::Chat.new
+
+# Set messages from an array of hashes
+chat.messages = [
+  { role: "system", content: "You are a helpful assistant." },
+  { role: "user", content: "Hello!" },
+  { role: "assistant", content: "How can I help you today?" }
+]
+
+# Set messages from ActiveRecord models
+chat.messages = Message.where(conversation_id: 123)
+
+# With images
+chat.messages = [
+  { role: "system", content: "You are a helpful assistant." },
+  { role: "user", content: "What's in this image?", image: "path/to/image.jpg" },
+  { role: "assistant", content: "I see a cat in the image." }
+]
+
+# With multiple images
+chat.messages = [
+  { role: "user", content: "Compare these images", images: ["image1.jpg", "image2.jpg"] }
+]
+```
+
+### Custom attribute mappings
+
+If your database columns or object attributes have different names, you can configure custom mappings:
+
+```ruby
+# Configure custom attribute mappings
+chat = AI::Chat.new
+chat.configure_attributes(
+  role: :message_type,       # Instead of "role"
+  content: :message_body,    # Instead of "content" 
+  images: :attachments,      # For retrieving associated images
+  image_url: :url            # Column on the image model that contains the URL/path
+)
+
+# Now works with custom column names
+chat.messages = CustomMessage.where(conversation_id: 123)
+```
 
 ## Testing with Real API Calls
 
@@ -212,7 +257,6 @@ Setting to `nil` disables the reasoning parameter.
 
 ## TODOs
 
-- Add the ability to set all messages at once, ideally with an ActiveRecord Relation.
 - Add a way to access the whole API response body (rather than just the message content).
 
 ## Contributing

lib/ai/chat.rb

@@ -4,8 +4,8 @@
 
 module AI
   class Chat
-    attr_accessor :messages, :schema, :model
-    attr_reader :reasoning_effort
+    attr_accessor :schema, :model
+    attr_reader :messages, :reasoning_effort, :attribute_mappings
 
     VALID_REASONING_EFFORTS = [:low, :medium, :high].freeze
 
@@ -14,6 +14,13 @@ def initialize(api_key: nil)
       @messages = []
       @model = "gpt-4.1-mini"
       @reasoning_effort = nil
+      @attribute_mappings = {
+        role: :role,
+        content: :content,
+        image: :image,
+        images: :images,
+        image_url: :image_url
+      }
     end
 
     def reasoning_effort=(value)
@@ -270,5 +277,78 @@ def process_image(obj)
         "data:#{mime_type};base64,#{base64_string}"
       end
     end
+
+    def messages=(new_messages)
+      # Reset the current messages array
+      @messages = []
+
+      # Process each message in the new_messages array/relation
+      new_messages.each do |message|
+        # Extract role and content using the configured attribute names
+        role = extract_attribute(message, @attribute_mappings[:role])
+        content = extract_attribute(message, @attribute_mappings[:content])
+
+        case role&.to_s
+        when "system"
+          system(content)
+        when "user"
+          # Handle images through various possible structures
+          if content.is_a?(Array)
+            # This is already a mixed content array with text and images
+            user(content)
+          else
+            # Extract images using configured attribute names
+            image = extract_attribute(message, @attribute_mappings[:image])
+            images = extract_attribute(message, @attribute_mappings[:images])
+
+            # For ActiveRecord associations that return collections
+            if images.nil? && message.respond_to?(@attribute_mappings[:images])
+              collection = message.send(@attribute_mappings[:images])
+              if collection.respond_to?(:each) && !collection.is_a?(String)
+                images = collection.map { |img| extract_attribute(img, @attribute_mappings[:image_url]) || img }
+              end
+            end
+
+            # Add the message with any found images
+            if image || (images && !images.empty?)
+              user(content, image: image, images: images)
+            else
+              user(content)
+            end
+          end
+        when "assistant"
+          assistant(content)
+        else
+          # For unknown roles, add directly but ensure symbols for keys
+          if message.is_a?(Hash)
+            @messages << message.transform_keys(&:to_sym)
+          else
+            # For ActiveRecord objects, convert to hash with symbol keys
+            hash = { role: role, content: content }
+            @messages << hash
+          end
+        end
+      end
+    end
+
+    # Configure attribute mappings
+    def configure_attributes(mappings = {})
+      mappings.each do |key, value|
+        @attribute_mappings[key.to_sym] = value.to_sym
+      end
+    end
+
+    # Helper method to extract an attribute from various object types
+    def extract_attribute(obj, attr_name)
+      if obj.respond_to?(attr_name)
+        # Method access (ActiveRecord)
+        obj.send(attr_name)
+      elsif obj.is_a?(Hash) && (obj.key?(attr_name) || obj.key?(attr_name.to_s))
+        # Hash access with symbol or string keys
+        obj[attr_name] || obj[attr_name.to_s]
+      else
+        nil
+      end
+    end
   end
 end

spec/ai/chat/messages_assignment_spec.rb

@@ -0,0 +1,180 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe AI::Chat, "messages assignment" do
+  let(:chat) { build(:chat) }
+  let(:test_image_path) { File.join(File.dirname(__FILE__), "../../fixtures/test1.jpg") }
+  let(:test_image_url) { "https://example.com/image.jpg" }
+
+  describe "#messages=" do
+    it "replaces existing messages with new ones" do
+      # Add some initial messages
+      chat.system("Initial system message")
+      chat.user("Initial user message")
+
+      # Replace with new messages
+      chat.messages = [
+        { role: "system", content: "New system message" },
+        { role: "user", content: "New user message" }
+      ]
+
+      expect(chat.messages.length).to eq(2)
+      expect(chat.messages[0][:role]).to eq("system")
+      expect(chat.messages[0][:content]).to eq("New system message")
+      expect(chat.messages[1][:role]).to eq("user")
+      expect(chat.messages[1][:content]).to eq("New user message")
+    end
+
+    it "works with string keys" do
+      chat.messages = [
+        { "role" => "system", "content" => "System message" },
+        { "role" => "user", "content" => "User message" }
+      ]
+
+      expect(chat.messages.length).to eq(2)
+      expect(chat.messages[0][:role]).to eq("system")
+      expect(chat.messages[0][:content]).to eq("System message")
+    end
+
+    it "handles messages with images" do
+      chat.messages = [
+        { role: "system", content: "System message" },
+        { role: "user", content: "User with image", image: test_image_path }
+      ]
+
+      expect(chat.messages.length).to eq(2)
+      expect(chat.messages[1][:role]).to eq("user")
+      expect(chat.messages[1][:content]).to be_an(Array)
+      expect(chat.messages[1][:content][0][:type]).to eq("input_text")
+      expect(chat.messages[1][:content][1][:type]).to eq("input_image")
+    end
+
+    it "handles messages with multiple images" do
+      chat.messages = [
+        { role: "user", content: "Multiple images", images: [test_image_path, test_image_url] }
+      ]
+
+      expect(chat.messages.length).to eq(1)
+      expect(chat.messages[0][:content]).to be_an(Array)
+      expect(chat.messages[0][:content].length).to eq(3) # text + 2 images
+    end
+
+    it "works with custom attribute mappings" do
+      chat.configure_attributes(
+        role: :message_type,
+        content: :message_body
+      )
+
+      chat.messages = [
+        { message_type: "system", message_body: "System with custom mapping" },
+        { message_type: "user", message_body: "User with custom mapping" }
+      ]
+
+      expect(chat.messages.length).to eq(2)
+      expect(chat.messages[0][:role]).to eq("system")
+      expect(chat.messages[0][:content]).to eq("System with custom mapping")
+    end
+
+    # Mock an ActiveRecord-like object
+    class MockMessage
+      attr_reader :message_type, :message_body
+
+      def initialize(type, body)
+        @message_type = type
+        @message_body = body
+      end
+    end
+
+    class MockMessageWithImage < MockMessage
+      attr_reader :image
+
+      def initialize(type, body, image)
+        super(type, body)
+        @image = image
+      end
+    end
+
+    class MockMessageWithImages < MockMessage
+      attr_reader :attachments
+
+      def initialize(type, body, attachments)
+        super(type, body)
+        @attachments = attachments
+      end
+    end
+
+    class MockImage
+      attr_reader :url
+
+      def initialize(url)
+        @url = url
+      end
+    end
+
+    it "works with object methods (ActiveRecord-like)" do
+      chat.configure_attributes(
+        role: :message_type,
+        content: :message_body
+      )
+
+      mock_messages = [
+        MockMessage.new("system", "System from object"),
+        MockMessage.new("user", "User from object")
+      ]
+
+      chat.messages = mock_messages
+
+      expect(chat.messages.length).to eq(2)
+      expect(chat.messages[0][:role]).to eq("system")
+      expect(chat.messages[0][:content]).to eq("System from object")
+    end
+
+    it "handles objects with image attributes" do
+      chat.configure_attributes(
+        role: :message_type,
+        content: :message_body,
+        image: :image
+      )
+
+      mock_messages = [
+        MockMessage.new("system", "System message"),
+        MockMessageWithImage.new("user", "Message with image", test_image_path)
+      ]
+
+      chat.messages = mock_messages
+
+      expect(chat.messages.length).to eq(2)
+      expect(chat.messages[1][:role]).to eq("user")
+      expect(chat.messages[1][:content]).to be_an(Array)
+      expect(chat.messages[1][:content][0][:type]).to eq("input_text")
+      expect(chat.messages[1][:content][1][:type]).to eq("input_image")
+    end
+
+    it "handles objects with image collections" do
+      chat.configure_attributes(
+        role: :message_type,
+        content: :message_body,
+        images: :attachments,
+        image_url: :url
+      )
+
+      mock_images = [
+        MockImage.new(test_image_path),
+        MockImage.new(test_image_url)
+      ]
+
+      mock_messages = [
+        MockMessage.new("system", "System message"),
+        MockMessageWithImages.new("user", "Message with images", mock_images)
+      ]
+
+      chat.messages = mock_messages
+
+      expect(chat.messages.length).to eq(2)
+      expect(chat.messages[1][:role]).to eq("user")
+      expect(chat.messages[1][:content]).to be_an(Array)
+      expect(chat.messages[1][:content].length).to eq(3) # text + 2 images
+    end
+  end
+end