chipping away at inline uploads

SamSaffron · SamSaffron · commit 8ce3cf1ef3b9 · 2025-03-27T15:08:29.000+11:00
diff --git a/lib/completions/prompt_messages_builder.rb b/lib/completions/prompt_messages_builder.rb
@@ -259,9 +259,9 @@ def push(type:, content:, name: nil, upload_ids: nil, id: nil, thinking: nil)
         end
         raise ArgumentError, "upload_ids must be an array" if upload_ids && !upload_ids.is_a?(Array)
 
+        content = [content, *upload_ids.map { |upload_id| { upload_id: upload_id } }] if upload_ids
         message = { type: type, content: content }
         message[:name] = name.to_s if name
-        message[:upload_ids] = upload_ids if upload_ids
         message[:id] = id.to_s if id
         if thinking
           message[:thinking] = thinking["thinking"] if thinking["thinking"]
@@ -331,16 +331,13 @@ def topic_array
 
       def chat_array(limit:)
         if @raw_messages.length > 1
-          buffer =
-            +"You are replying inside a Discourse chat channel. Here is a summary of the conversation so far:\n{{{"
-
-          upload_ids = []
+          buffer = [
+            +"You are replying inside a Discourse chat channel. Here is a summary of the conversation so far:\n{{{",
+          ]
 
           @raw_messages[0..-2].each do |message|
             buffer << "\n"
 
-            upload_ids.concat(message[:upload_ids]) if message[:upload_ids].present?
-
             if message[:type] == :user
               buffer << "#{message[:name] || "User"}: "
             else
@@ -357,15 +354,41 @@ def chat_array(limit:)
         end
 
         last_message = @raw_messages[-1]
-        buffer << "#{last_message[:name] || "User"}: #{last_message[:content]} "
+        buffer << "#{last_message[:name] || "User"}: "
+        buffer << last_message[:content]
+
+        buffer = compress_messages_buffer(buffer.flatten, max_uploads: MAX_CHAT_UPLOADS)
 
         message = { type: :user, content: buffer }
-        upload_ids.concat(last_message[:upload_ids]) if last_message[:upload_ids].present?
+        [message]
+      end
 
-        message[:upload_ids] = upload_ids[-MAX_CHAT_UPLOADS..-1] ||
-          upload_ids if upload_ids.present?
+      # caps uploads to maximum uploads allowed in message stream
+      # and concats string elements
+      def compress_messages_buffer(buffer, max_uploads:)
+        compressed = []
+        current_text = +""
+        upload_count = 0
+
+        buffer.each do |item|
+          if item.is_a?(String)
+            current_text << item
+          elsif item.is_a?(Hash)
+            compressed << current_text if current_text.present?
+            compressed << item
+            current_text = +""
+            upload_count += 1
+          end
+        end
 
-        [message]
+        compressed << current_text if current_text.present?
+
+        if upload_count > max_uploads
+          counter = max_uploads - upload_count
+          compressed.delete_if { |item| item.is_a?(Hash) && (counter += 1) > 0 }
+        end
+
+        compressed
       end
     end
   end
diff --git a/spec/lib/completions/prompt_messages_builder_spec.rb b/spec/lib/completions/prompt_messages_builder_spec.rb
@@ -2,6 +2,9 @@
 
 describe DiscourseAi::Completions::PromptMessagesBuilder do
   let(:builder) { DiscourseAi::Completions::PromptMessagesBuilder.new }
+  fab!(:user)
+  fab!(:bot_user) { Fabricate(:user) }
+  fab!(:other_user) { Fabricate(:user) }
 
   it "should allow merging user messages" do
     builder.push(type: :user, content: "Hello", name: "Alice")
@@ -14,7 +17,7 @@
     builder.push(type: :user, content: "Hello", name: "Alice", upload_ids: [1, 2])
 
     expect(builder.to_a).to eq(
-      [{ type: :user, name: "Alice", content: "Hello", upload_ids: [1, 2] }],
+      [{ type: :user, content: ["Hello", { upload_id: 1 }, { upload_id: 2 }], name: "Alice" }],
     )
   end
 
@@ -64,4 +67,265 @@
     expect(content).to include("Alice")
     expect(content).to include("How do I solve this")
   end
+
+  describe ".messages_from_chat" do
+    fab!(:dm_channel) { Fabricate(:direct_message_channel, users: [user, bot_user]) }
+    fab!(:dm_message1) do
+      Fabricate(:chat_message, chat_channel: dm_channel, user: user, message: "Hello bot")
+    end
+    fab!(:dm_message2) do
+      Fabricate(:chat_message, chat_channel: dm_channel, user: bot_user, message: "Hello human")
+    end
+    fab!(:dm_message3) do
+      Fabricate(:chat_message, chat_channel: dm_channel, user: user, message: "How are you?")
+    end
+
+    fab!(:public_channel) { Fabricate(:category_channel) }
+    fab!(:public_message1) do
+      Fabricate(:chat_message, chat_channel: public_channel, user: user, message: "Hello everyone")
+    end
+    fab!(:public_message2) do
+      Fabricate(:chat_message, chat_channel: public_channel, user: bot_user, message: "Hi there")
+    end
+
+    fab!(:thread_original) do
+      Fabricate(:chat_message, chat_channel: public_channel, user: user, message: "Thread starter")
+    end
+    fab!(:thread) do
+      Fabricate(:chat_thread, channel: public_channel, original_message: thread_original)
+    end
+    fab!(:thread_reply1) do
+      Fabricate(
+        :chat_message,
+        chat_channel: public_channel,
+        user: other_user,
+        message: "Thread reply",
+        thread: thread,
+      )
+    end
+
+    fab!(:upload) { Fabricate(:upload, user: user) }
+    fab!(:message_with_upload) do
+      Fabricate(
+        :chat_message,
+        chat_channel: dm_channel,
+        user: user,
+        message: "Check this image",
+        upload_ids: [upload.id],
+      )
+    end
+
+    it "processes messages from direct message channels" do
+      context =
+        described_class.messages_from_chat(
+          dm_message3,
+          channel: dm_channel,
+          context_post_ids: nil,
+          max_messages: 10,
+          include_uploads: false,
+          bot_user_ids: [bot_user.id],
+          instruction_message: nil,
+        )
+
+      # this is all we got cause it is assuming threading
+      expect(context).to eq([{ type: :user, content: "How are you?", name: user.username }])
+    end
+
+    it "includes uploads when include_uploads is true" do
+      message_with_upload.reload
+      expect(message_with_upload.uploads).to include(upload)
+
+      context =
+        described_class.messages_from_chat(
+          message_with_upload,
+          channel: dm_channel,
+          context_post_ids: nil,
+          max_messages: 10,
+          include_uploads: true,
+          bot_user_ids: [bot_user.id],
+          instruction_message: nil,
+        )
+
+      # Find the message with upload
+      message = context.find { |m| m[:content] == ["Check this image", { upload_id: upload.id }] }
+      expect(message).to be_present
+    end
+
+    it "doesn't include uploads when include_uploads is false" do
+      # Make sure the upload is associated with the message
+      message_with_upload.reload
+      expect(message_with_upload.uploads).to include(upload)
+
+      context =
+        described_class.messages_from_chat(
+          message_with_upload,
+          channel: dm_channel,
+          context_post_ids: nil,
+          max_messages: 10,
+          include_uploads: false,
+          bot_user_ids: [bot_user.id],
+          instruction_message: nil,
+        )
+
+      # Find the message with upload
+      message = context.find { |m| m[:content] == "Check this image" }
+      expect(message).to be_present
+      expect(message[:upload_ids]).to be_nil
+    end
+
+    it "properly handles uploads in public channels with multiple users" do
+      _first_message =
+        Fabricate(:chat_message, chat_channel: public_channel, user: user, message: "First message")
+
+      _message_with_upload =
+        Fabricate(
+          :chat_message,
+          chat_channel: public_channel,
+          user: other_user,
+          message: "Message with image",
+          upload_ids: [upload.id],
+        )
+
+      last_message =
+        Fabricate(:chat_message, chat_channel: public_channel, user: user, message: "Final message")
+
+      context =
+        described_class.messages_from_chat(
+          last_message,
+          channel: public_channel,
+          context_post_ids: nil,
+          max_messages: 3,
+          include_uploads: true,
+          bot_user_ids: [bot_user.id],
+          instruction_message: nil,
+        )
+
+      expect(context.length).to eq(1)
+      content = context.first[:content]
+
+      expect(content.length).to eq(3)
+      expect(content[0]).to include("First message")
+      expect(content[0]).to include("Message with image")
+      expect(content[1]).to include({ upload_id: upload.id })
+      expect(content[2]).to include("Final message")
+    end
+  end
+
+  describe ".messages_from_post" do
+    fab!(:pm) do
+      Fabricate(
+        :private_message_topic,
+        title: "This is my special PM",
+        user: user,
+        topic_allowed_users: [
+          Fabricate.build(:topic_allowed_user, user: user),
+          Fabricate.build(:topic_allowed_user, user: bot_user),
+        ],
+      )
+    end
+    fab!(:first_post) do
+      Fabricate(:post, topic: pm, user: user, post_number: 1, raw: "This is a reply by the user")
+    end
+    fab!(:second_post) do
+      Fabricate(:post, topic: pm, user: bot_user, post_number: 2, raw: "This is a bot reply")
+    end
+    fab!(:third_post) do
+      Fabricate(
+        :post,
+        topic: pm,
+        user: user,
+        post_number: 3,
+        raw: "This is a second reply by the user",
+      )
+    end
+
+    context "with limited context" do
+      it "respects max_context_posts" do
+        context =
+          described_class.messages_from_post(
+            third_post,
+            max_posts: 1,
+            bot_usernames: [bot_user.username],
+            include_uploads: false,
+          )
+
+        expect(context).to contain_exactly(
+          *[{ type: :user, id: user.username, content: third_post.raw }],
+        )
+      end
+    end
+
+    it "includes previous posts ordered by post_number" do
+      context =
+        described_class.messages_from_post(
+          third_post,
+          max_posts: 10,
+          bot_usernames: [bot_user.username],
+          include_uploads: false,
+        )
+
+      expect(context).to eq(
+        [
+          { type: :user, content: "This is a reply by the user", id: user.username },
+          { type: :model, content: "This is a bot reply" },
+          { type: :user, content: "This is a second reply by the user", id: user.username },
+        ],
+      )
+    end
+
+    it "only include regular posts" do
+      first_post.update!(post_type: Post.types[:whisper])
+
+      context =
+        described_class.messages_from_post(
+          third_post,
+          max_posts: 10,
+          bot_usernames: [bot_user.username],
+          include_uploads: false,
+        )
+
+      # skips leading model reply which makes no sense cause first post was whisper
+      expect(context).to eq(
+        [{ type: :user, content: "This is a second reply by the user", id: user.username }],
+      )
+    end
+
+    context "with custom prompts" do
+      it "When post custom prompt is present, we use that instead of the post content" do
+        custom_prompt = [
+          [
+            { name: "time", arguments: { name: "time", timezone: "Buenos Aires" } }.to_json,
+            "time",
+            "tool_call",
+          ],
+          [
+            { args: { timezone: "Buenos Aires" }, time: "2023-12-14 17:24:00 -0300" }.to_json,
+            "time",
+            "tool",
+          ],
+          ["I replied to the time command", bot_user.username],
+        ]
+
+        PostCustomPrompt.create!(post: second_post, custom_prompt: custom_prompt)
+
+        context =
+          described_class.messages_from_post(
+            third_post,
+            max_posts: 10,
+            bot_usernames: [bot_user.username],
+            include_uploads: false,
+          )
+
+        expect(context).to eq(
+          [
+            { type: :user, content: "This is a reply by the user", id: user.username },
+            { type: :tool_call, content: custom_prompt.first.first, id: "time" },
+            { type: :tool, id: "time", content: custom_prompt.second.first },
+            { type: :model, content: custom_prompt.third.first },
+            { type: :user, content: "This is a second reply by the user", id: user.username },
+          ],
+        )
+      end
+    end
+  end
 end
diff --git a/spec/lib/modules/ai_bot/playground_spec.rb b/spec/lib/modules/ai_bot/playground_spec.rb
