FEATURE: scale up result count for search depending on model (#346)

We were limiting to 20 results unconditionally cause we had to make
sure search always fit in an 8k context window.

Models such as GPT 3.5 Turbo (16k) and GPT 4 Turbo / Claude 2.1 (over 150k)
allow us to return a lot more results.

This means we have a much richer understanding cause context is far
larger.

This also allows a persona to tweak this number, in some cases admin
may want to be conservative and save on tokens by limiting results

This also tweaks the `limit` param which GPT-4 liked to set to tell
model only to use it when it needs to (and describes default behavior)

assets/javascripts/initializers/ai-bot-replies.js

@@ -8,7 +8,7 @@ import { registerWidgetShim } from "discourse/widgets/render-glimmer";
 import { composeAiBotMessage } from "discourse/plugins/discourse-ai/discourse/lib/ai-bot-helper";
 
 function isGPTBot(user) {
-  return user && [-110, -111, -112].includes(user.id);
+  return user && [-110, -111, -112, -113].includes(user.id);
 }
 
 function attachHeaderIcon(api) {

config/locales/server.en.yml

@@ -155,6 +155,9 @@ en:
       searching: "Searching for: '%{query}'"
       command_options:
         search:
+          max_results:
+            name: "Maximum number of results"
+            description: "Maximum number of results to include in the search - if empty default rules will be used and count will be scaled depending on model used. Highest value is 100."
           base_query:
             name: "Base Search Query"
             description: "Base query to use when searching. Example: '#urgent' will prepend '#urgent' to the search query and only include topics with the urgent category or tag."

lib/ai_bot/commands/search_command.rb

@@ -12,7 +12,7 @@ module DiscourseAi::AiBot::Commands
       end
 
       def options
-        [option(:base_query, type: :string)]
+        [option(:base_query, type: :string), option(:max_results, type: :integer)]
       end
 
       def parameters
@@ -38,7 +38,7 @@ module DiscourseAi::AiBot::Commands
           Parameter.new(
             name: "limit",
             description:
-              "limit number of results returned (generally prefer to just keep to default)",
+              "Number of results to return. Defaults to maximum number of results. Only set if absolutely necessary",
             type: "integer",
           ),
           Parameter.new(
@@ -98,9 +98,27 @@ module DiscourseAi::AiBot::Commands
       }
     end
 
-    MAX_RESULTS = 20
     MIN_SEMANTIC_RESULTS = 5
 
+    def max_semantic_results
+      max_results / 4
+    end
+
+    def max_results
+      return 20 if !bot
+
+      max_results = persona_options[:max_results].to_i
+      return [max_results, 100].min if max_results > 0
+
+      if bot.prompt_limit(allow_commands: false) > 30_000
+        60
+      elsif bot.prompt_limit(allow_commands: false) > 10_000
+        40
+      else
+        20
+      end
+    end
+
     def process(**search_args)
       limit = nil
 
@@ -135,14 +153,14 @@ module DiscourseAi::AiBot::Commands
         )
 
       # let's be frugal with tokens, 50 results is too much and stuff gets cut off
-      limit ||= MAX_RESULTS
+      limit ||= max_results
-      limit = MAX_RESULTS if limit > MAX_RESULTS
+      limit = max_results if limit > max_results
 
       should_try_semantic_search = SiteSetting.ai_embeddings_semantic_search_enabled
-      should_try_semantic_search &&= (limit == MAX_RESULTS)
+      should_try_semantic_search &&= (limit == max_results)
       should_try_semantic_search &&= (search_args[:search_query].present?)
 
-      limit = limit - MIN_SEMANTIC_RESULTS if should_try_semantic_search
+      limit = limit - max_semantic_results if should_try_semantic_search
 
       posts = results&.posts || []
       posts = posts[0..limit - 1]
@@ -169,7 +187,7 @@ module DiscourseAi::AiBot::Commands
             topic_ids << post.topic_id
             posts << post
 
-            break if posts.length >= MAX_RESULTS
+            break if posts.length >= max_results
           end
         end
       end

spec/lib/modules/ai_bot/commands/search_command_spec.rb

@@ -27,11 +27,15 @@ RSpec.describe DiscourseAi::AiBot::Commands::SearchCommand do
   before { SiteSetting.ai_bot_enabled = true }
 
   it "can properly list options" do
-    options = described_class.options
+    options = described_class.options.sort_by(&:name)
-    expect(options.length).to eq(1)
+    expect(options.length).to eq(2)
     expect(options.first.name.to_s).to eq("base_query")
     expect(options.first.localized_name).not_to include("Translation missing:")
     expect(options.first.localized_description).not_to include("Translation missing:")
+
+    expect(options.second.name.to_s).to eq("max_results")
+    expect(options.second.localized_name).not_to include("Translation missing:")
+    expect(options.second.localized_description).not_to include("Translation missing:")
   end
 
   describe "#process" do
@@ -137,6 +141,35 @@ RSpec.describe DiscourseAi::AiBot::Commands::SearchCommand do
       expect(tags).to eq("funny, sad")
     end
 
+    it "scales results to number of tokens" do
+      SiteSetting.ai_bot_enabled_chat_bots = "gpt-3.5-turbo|gpt-4|claude-2"
+
+      post1 = Fabricate(:post)
+
+      gpt_3_5_turbo =
+        DiscourseAi::AiBot::Bot.as(User.find(DiscourseAi::AiBot::EntryPoint::GPT3_5_TURBO_ID))
+      gpt4 = DiscourseAi::AiBot::Bot.as(User.find(DiscourseAi::AiBot::EntryPoint::GPT4_ID))
+      claude = DiscourseAi::AiBot::Bot.as(User.find(DiscourseAi::AiBot::EntryPoint::CLAUDE_V2_ID))
+
+      expect(described_class.new(bot: claude, post: post1, args: nil).max_results).to eq(60)
+      expect(described_class.new(bot: gpt_3_5_turbo, post: post1, args: nil).max_results).to eq(40)
+      expect(described_class.new(bot: gpt4, post: post1, args: nil).max_results).to eq(20)
+
+      persona =
+        Fabricate(
+          :ai_persona,
+          commands: [["SearchCommand", { "max_results" => 6 }]],
+          enabled: true,
+          allowed_group_ids: [Group::AUTO_GROUPS[:trust_level_0]],
+        )
+
+      Group.refresh_automatic_groups!
+
+      custom_bot = DiscourseAi::AiBot::Bot.as(bot_user, persona_id: persona.id, user: admin)
+
+      expect(described_class.new(bot: custom_bot, post: post1, args: nil).max_results).to eq(6)
+    end
+
     it "can handle limits" do
       post1 = Fabricate(:post, topic: topic_with_tags)
       _post2 = Fabricate(:post, user: post1.user)

spec/requests/admin/ai_personas_controller_spec.rb

@@ -46,6 +46,12 @@ RSpec.describe DiscourseAi::Admin::AiPersonasController do
             "description" =>
               I18n.t("discourse_ai.ai_bot.command_options.search.base_query.description"),
           },
+          "max_results" => {
+            "type" => "integer",
+            "name" => I18n.t("discourse_ai.ai_bot.command_options.search.max_results.name"),
+            "description" =>
+              I18n.t("discourse_ai.ai_bot.command_options.search.max_results.description"),
+          },
         },
       )