diff --git a/assets/javascripts/initializers/ai-bot-replies.js b/assets/javascripts/initializers/ai-bot-replies.js index aa7c39bd..e136963e 100644 --- a/assets/javascripts/initializers/ai-bot-replies.js +++ b/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) { diff --git a/config/locales/server.en.yml b/config/locales/server.en.yml index 7bce1c11..390bee93 100644 --- a/config/locales/server.en.yml +++ b/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." diff --git a/lib/ai_bot/commands/search_command.rb b/lib/ai_bot/commands/search_command.rb index 7da2012b..9cd535fc 100644 --- a/lib/ai_bot/commands/search_command.rb +++ b/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 if limit > MAX_RESULTS + 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 diff --git a/spec/lib/modules/ai_bot/commands/search_command_spec.rb b/spec/lib/modules/ai_bot/commands/search_command_spec.rb index 7026aa22..dc00b3e3 100644 --- a/spec/lib/modules/ai_bot/commands/search_command_spec.rb +++ b/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 - expect(options.length).to eq(1) + options = described_class.options.sort_by(&:name) + 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) diff --git a/spec/requests/admin/ai_personas_controller_spec.rb b/spec/requests/admin/ai_personas_controller_spec.rb index 31d21163..1f75481f 100644 --- a/spec/requests/admin/ai_personas_controller_spec.rb +++ b/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"), + }, }, )