DEV: refactor bot internals

This introduces a proper object for bot context, this makes
it simpler to improve context management as we go cause we
have a nice object to work with

Starts refactoring allowing for a single message to have
multiple uploads throughout

Author: SamSaffron

lib/ai_bot/bot.rb

@@ -75,17 +75,17 @@ def get_updated_title(conversation_context, post, user)
       def force_tool_if_needed(prompt, context)
         return if prompt.tool_choice == :none
 
-        context[:chosen_tools] ||= []
+        context.chosen_tools ||= []
         forced_tools = persona.force_tool_use.map { |tool| tool.name }
-        force_tool = forced_tools.find { |name| !context[:chosen_tools].include?(name) }
+        force_tool = forced_tools.find { |name| !context.chosen_tools.include?(name) }
 
         if force_tool && persona.forced_tool_count > 0
           user_turns = prompt.messages.select { |m| m[:type] == :user }.length
           force_tool = false if user_turns > persona.forced_tool_count
         end
 
         if force_tool
-          context[:chosen_tools] << force_tool
+          context.chosen_tools << force_tool
           prompt.tool_choice = force_tool
         else
           prompt.tool_choice = nil
@@ -100,7 +100,7 @@ def reply(context, &update_blk)
         ongoing_chain = true
         raw_context = []
 
-        user = context[:user]
+        user = context.user
 
         llm_kwargs = { user: user }
         llm_kwargs[:temperature] = persona.temperature if persona.temperature
@@ -297,7 +297,7 @@ def process_tool(
       end
 
       def invoke_tool(tool, llm, cancel, context, &update_blk)
-        show_placeholder = !context[:skip_tool_details] && !tool.class.allow_partial_tool_calls?
+        show_placeholder = !context.skip_tool_details && !tool.class.allow_partial_tool_calls?
 
         update_blk.call("", cancel, build_placeholder(tool.summary, "")) if show_placeholder
 

lib/ai_bot/bot_context.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module AiBot
+    class BotContext
+      attr_accessor :messages,
+                    :topic_id,
+                    :post_id,
+                    :private_message,
+                    :custom_instructions,
+                    :user,
+                    :skip_tool_details,
+                    :participants,
+                    :chosen_tools
+
+      def initialize(
+        post: nil,
+        participants: nil,
+        user: nil,
+        skip_tool_details: nil,
+        messages:,
+        custom_instructions: nil
+      )
+        @participants = participants
+        @user = user
+        @skip_tool_details = skip_tool_details
+        @messages = messages
+        @custom_instructions = custom_instructions
+
+        if post
+          @post_id = post.id
+          @topic_id = post.topic_id
+          @private_message = post.topic.private_message?
+          @participants = post.topic.allowed_users.map(&:username).join(", ") if @private_message
+          @user = post.user
+        end
+      end
+
+      # these are strings that can be safely interpolated into templates
+      TEMPLATE_PARAMS = %w[time site_url site_title site_description]
+
+      def lookup_template_param(key)
+        public_send(key.to_sym) if TEMPLATE_PARAMS.include?(key)
+      end
+
+      def time
+        @time ||= Time.zone.now
+      end
+
+      def site_url
+        Discourse.base_url
+      end
+
+      def site_title
+        SiteSetting.title
+      end
+
+      def site_description
+        SiteSetting.site_description
+      end
+
+      def private_message?
+        @private_message
+      end
+    end
+  end
+end

lib/ai_bot/personas/persona.rb

@@ -163,7 +163,7 @@ def available_tools
         def craft_prompt(context, llm: nil)
           system_insts =
             system_prompt.gsub(/\{(\w+)\}/) do |match|
-              found = context[match[1..-2].to_sym]
+              found = context.lookup_template_param(match[1..-2])
               found.nil? ? match : found.to_s
             end
 
@@ -180,26 +180,26 @@ def craft_prompt(context, llm: nil)
               )
           end
 
-          if context[:custom_instructions].present?
+          if context.custom_instructions.present?
             prompt_insts << "\n"
-            prompt_insts << context[:custom_instructions]
+            prompt_insts << context.custom_instructions
           end
 
           fragments_guidance =
             rag_fragments_prompt(
-              context[:conversation_context].to_a,
+              context.messages,
               llm: question_consolidator_llm,
-              user: context[:user],
+              user: context.user,
             )&.strip
 
           prompt_insts << fragments_guidance if fragments_guidance.present?
 
           prompt =
             DiscourseAi::Completions::Prompt.new(
               prompt_insts,
-              messages: context[:conversation_context].to_a,
-              topic_id: context[:topic_id],
-              post_id: context[:post_id],
+              messages: context.messages,
+              topic_id: context.topic_id,
+              post_id: context.post_id,
             )
 
           prompt.max_pixels = self.class.vision_max_pixels if self.class.vision_enabled

lib/ai_bot/playground.rb

@@ -227,7 +227,7 @@ def update_playground_with(post)
         schedule_bot_reply(post) if can_attach?(post)
       end
 
-      def conversation_context(post, style: nil)
+      def post_prompt_messages(post, style: nil)
         # Pay attention to the `post_number <= ?` here.
         # We want to inject the last post as context because they are translated differently.
 
@@ -307,10 +307,10 @@ def conversation_context(post, style: nil)
       end
 
       def title_playground(post, user)
-        context = conversation_context(post)
+        messages = post_prompt_messages(post)
 
         bot
-          .get_updated_title(context, post, user)
+          .get_updated_title(messages, post, user)
           .tap do |new_title|
             PostRevisor.new(post.topic.first_post, post.topic).revise!(
               bot.bot_user,
@@ -326,7 +326,7 @@ def title_playground(post, user)
         )
       end
 
-      def chat_context(message, channel, persona_user, context_post_ids)
+      def chat_prompt_messages(message, channel, persona_user, context_post_ids)
         has_vision = bot.persona.class.vision_enabled
         include_thread_titles = !channel.direct_message_channel? && !message.thread_id
 
@@ -411,9 +411,9 @@ def reply_to_chat_message(message, channel, context_post_ids)
         context_post_ids = nil if !channel.direct_message_channel?
 
         context =
-          get_context(
-            participants: participants.join(", "),
-            conversation_context: chat_context(message, channel, persona_user, context_post_ids),
+          BotContext.new(
+            participants: participants,
+            messages: chat_prompt_messages(message, channel, persona_user, context_post_ids),
             user: message.user,
             skip_tool_details: true,
           )
@@ -460,22 +460,6 @@ def reply_to_chat_message(message, channel, context_post_ids)
         reply
       end
 
-      def get_context(participants:, conversation_context:, user:, skip_tool_details: nil)
-        result = {
-          site_url: Discourse.base_url,
-          site_title: SiteSetting.title,
-          site_description: SiteSetting.site_description,
-          time: Time.zone.now,
-          participants: participants,
-          conversation_context: conversation_context,
-          user: user,
-        }
-
-        result[:skip_tool_details] = true if skip_tool_details
-
-        result
-      end
-
       def reply_to(
         post,
         custom_instructions: nil,
@@ -510,15 +494,11 @@ def reply_to(
           )
 
         context =
-          get_context(
-            participants: post.topic.allowed_users.map(&:username).join(", "),
-            conversation_context: conversation_context(post, style: context_style),
-            user: post.user,
+          BotContext.new(
+            post: post,
+            custom_instructions: custom_instructions,
+            messages: post_prompt_messages(post, style: context_style),
           )
-        context[:post_id] = post.id
-        context[:topic_id] = post.topic_id
-        context[:private_message] = post.topic.private_message?
-        context[:custom_instructions] = custom_instructions
 
         reply_user = bot.bot_user
         if bot.persona.class.respond_to?(:user_id)
@@ -562,7 +542,7 @@ def reply_to(
           Discourse.redis.setex(redis_stream_key, 60, 1)
         end
 
-        context[:skip_tool_details] ||= !bot.persona.class.tool_details
+        context.skip_tool_details ||= !bot.persona.class.tool_details
 
         post_streamer = PostStreamer.new(delay: Rails.env.test? ? 0 : 0.5) if stream_reply
 

lib/ai_bot/tool_runner.rb

@@ -457,7 +457,7 @@ def attach_upload(mini_racer_context)
                     UploadCreator.new(
                       file,
                       filename,
-                      for_private_message: @context[:private_message],
+                      for_private_message: @context.private_message,
                     ).create_for(@bot_user.id)
 
                   { id: upload.id, short_url: upload.short_url, url: upload.url }

lib/ai_bot/tools/dall_e.rb

@@ -111,7 +111,7 @@ def invoke
                     UploadCreator.new(
                       file,
                       "image.png",
-                      for_private_message: context[:private_message],
+                      for_private_message: context.private_message?,
                     ).create_for(bot_user.id),
                 }
               end

lib/completions/prompt.rb

@@ -115,13 +115,28 @@ def has_tools?
         tools.present?
       end
 
-      # helper method to get base64 encoded uploads
-      # at the correct dimentions
+      # TODO migrate to new API
       def encoded_uploads(message)
         return [] if message[:upload_ids].blank?
         UploadEncoder.encode(upload_ids: message[:upload_ids], max_pixels: max_pixels)
       end
 
+      def encode_upload(upload_id)
+        UploadEncoder.encode(upload_ids: [upload_id], max_pixels: max_pixels).first
+      end
+
+      def content_with_encoded_uploads(content)
+        return [content] unless content.is_a?(Array)
+
+        content.map do |c|
+          if c.is_a?(Hash) && c.key?(:upload_id)
+            encode_upload(c[:upload_id])
+          else
+            c
+          end
+        end
+      end
+
       def ==(other)
         return false unless other.is_a?(Prompt)
         messages == other.messages && tools == other.tools && topic_id == other.topic_id &&
@@ -166,8 +181,17 @@ def validate_message(message)
         if message[:upload_ids].present? && message[:type] != :user
           raise ArgumentError, "upload_ids are only supported for users"
         end
-
-        raise ArgumentError, "message content must be a string" if !message[:content].is_a?(String)
+        if message[:content].is_a?(Array)
+          message[:content].each do |content|
+            if !content.is_a?(String) && !(content.is_a?(Hash) && content.keys == [:upload_id])
+              raise ArgumentError, "Array message content must be a string or {upload_id: ...} "
+            end
+          end
+        else
+          if !message[:content].is_a?(String)
+            raise ArgumentError, "Message content must be a string or an array"
+          end
+        end
       end
 
       def validate_turn(last_turn, new_turn)

spec/lib/completions/prompt_spec.rb

@@ -25,6 +25,23 @@
   end
 
   describe "image support" do
+    it "allows adding uploads inline in messages" do
+      upload = UploadCreator.new(image100x100, "image.jpg").create_for(Discourse.system_user.id)
+
+      prompt.max_pixels = 300
+      prompt.push(
+        type: :user,
+        content: ["this is an image", { upload_id: upload.id }, "this was an image"],
+      )
+
+      encoded = prompt.content_with_encoded_uploads(prompt.messages.last[:content])
+
+      expect(encoded.length).to eq(3)
+      expect(encoded[0]).to eq("this is an image")
+      expect(encoded[1][:mime_type]).to eq("image/jpeg")
+      expect(encoded[2]).to eq("this was an image")
+    end
+
     it "allows adding uploads to messages" do
       upload = UploadCreator.new(image100x100, "image.jpg").create_for(Discourse.system_user.id)
 

spec/lib/modules/ai_bot/playground_spec.rb

@@ -1155,7 +1155,7 @@
     end
   end
 
-  describe "#conversation_context" do
+  describe "#post_prompt_messages" do
     context "with limited context" do
       before do
         @old_persona = playground.bot.persona
@@ -1166,7 +1166,7 @@
       after { playground.bot.persona = @old_persona }
 
       it "respects max_context_post" do
-        context = playground.conversation_context(third_post)
+        context = playground.post_prompt_messages(third_post)
 
         expect(context).to contain_exactly(
           *[{ type: :user, id: user.username, content: third_post.raw }],
@@ -1215,7 +1215,7 @@
 
         PostCustomPrompt.create!(post: second_post, custom_prompt: custom_prompt)
 
-        context = playground.conversation_context(third_post)
+        context = playground.post_prompt_messages(third_post)
 
         expect(context).to contain_exactly(
           *[