From 61890b667c06299841ae88946f84a112f00060e1 Mon Sep 17 00:00:00 2001 From: Sam Date: Fri, 10 May 2024 11:32:34 +1000 Subject: [PATCH] FEATURE: search command now support searching in context of user (#610) This optional feature allows search to be performed in the context of the user that executed it. By default we do not allow this behavior cause it means llm gets access to potentially secure data. --- .../ai-persona-command-option-editor.gjs | 55 ++++++++++++++----- config/locales/server.en.yml | 3 + lib/ai_bot/tools/search.rb | 26 ++++++--- lib/ai_bot/tools/tool.rb | 20 ++++--- spec/lib/modules/ai_bot/tools/search_spec.rb | 32 ++++++++++- .../admin/ai_personas_controller_spec.rb | 6 ++ 6 files changed, 110 insertions(+), 32 deletions(-) diff --git a/assets/javascripts/discourse/components/ai-persona-command-option-editor.gjs b/assets/javascripts/discourse/components/ai-persona-command-option-editor.gjs index e5dab0b5..7d14802c 100644 --- a/assets/javascripts/discourse/components/ai-persona-command-option-editor.gjs +++ b/assets/javascripts/discourse/components/ai-persona-command-option-editor.gjs @@ -1,17 +1,44 @@ +import Component from "@glimmer/component"; import { Input } from "@ember/component"; +import { on } from "@ember/modifier"; +import { action } from "@ember/object"; -const AiPersonaCommandOptionEditor = ; +export default class AiPersonaCommandOptionEditor extends Component { + get isBoolean() { + return this.args.option.type === "boolean"; + } -export default AiPersonaCommandOptionEditor; + get selectedValue() { + return this.args.option.value.value === "true"; + } + + @action + onCheckboxChange(event) { + this.args.option.value.value = event.target.checked ? "true" : "false"; + } + + +} diff --git a/config/locales/server.en.yml b/config/locales/server.en.yml index 002585db..d3438549 100644 --- a/config/locales/server.en.yml +++ b/config/locales/server.en.yml @@ -205,6 +205,9 @@ en: searching: "Searching for: '%{query}'" command_options: search: + search_private: + name: "Search Private" + description: "Include all topics user has access to in search results (by default only public topics are included)" 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." diff --git a/lib/ai_bot/tools/search.rb b/lib/ai_bot/tools/search.rb index d29f6dd7..3bc91bee 100644 --- a/lib/ai_bot/tools/search.rb +++ b/lib/ai_bot/tools/search.rb @@ -83,7 +83,11 @@ module DiscourseAi end def accepted_options - [option(:base_query, type: :string), option(:max_results, type: :integer)] + [ + option(:base_query, type: :string), + option(:max_results, type: :integer), + option(:search_private, type: :boolean), + ] end end @@ -106,12 +110,18 @@ module DiscourseAi search_string = "#{search_string} #{options[:base_query]}" end + safe_search_string = search_string.to_s + guardian = nil + + if options[:search_private] && context[:user] + guardian = Guardian.new(context[:user]) + else + guardian = Guardian.new + safe_search_string += " status:public" + end + results = - ::Search.execute( - search_string.to_s + " status:public", - search_type: :full_page, - guardian: Guardian.new(), - ) + ::Search.execute(safe_search_string, search_type: :full_page, guardian: guardian) # let's be frugal with tokens, 50 results is too much and stuff gets cut off max_results = calculate_max_results(llm) @@ -129,10 +139,10 @@ module DiscourseAi posts = posts[0..results_limit.to_i - 1] if should_try_semantic_search - semantic_search = DiscourseAi::Embeddings::SemanticSearch.new(Guardian.new()) + semantic_search = DiscourseAi::Embeddings::SemanticSearch.new(guardian) topic_ids = Set.new(posts.map(&:topic_id)) - search = ::Search.new(search_string, guardian: Guardian.new) + search = ::Search.new(search_string, guardian: guardian) results = nil begin diff --git a/lib/ai_bot/tools/tool.rb b/lib/ai_bot/tools/tool.rb index 9ccd2a1e..71c92eb8 100644 --- a/lib/ai_bot/tools/tool.rb +++ b/lib/ai_bot/tools/tool.rb @@ -66,14 +66,20 @@ module DiscourseAi end def options - self - .class - .accepted_options - .reduce(HashWithIndifferentAccess.new) do |memo, option| - val = @persona_options[option.name] - memo[option.name] = val if val - memo + result = HashWithIndifferentAccess.new + self.class.accepted_options.each do |option| + val = @persona_options[option.name] + if val + case option.type + when :boolean + val = val == "true" + when :integer + val = val.to_i + end + result[option.name] = val end + end + result end def chain_next_response? diff --git a/spec/lib/modules/ai_bot/tools/search_spec.rb b/spec/lib/modules/ai_bot/tools/search_spec.rb index a33f2a14..2c32146c 100644 --- a/spec/lib/modules/ai_bot/tools/search_spec.rb +++ b/spec/lib/modules/ai_bot/tools/search_spec.rb @@ -29,13 +29,29 @@ RSpec.describe DiscourseAi::AiBot::Tools::Search do Fabricate(:topic, category: category, tags: [tag_funny, tag_sad, tag_hidden]) end + fab!(:user) + fab!(:group) + fab!(:private_category) do + c = Fabricate(:category_with_definition) + c.set_permissions(group => :readonly) + c.save + c + end + before { SiteSetting.ai_bot_enabled = true } describe "#invoke" do it "can retrieve options from persona correctly" do - persona_options = { "base_query" => "#funny" } + persona_options = { + "base_query" => "#funny", + "search_private" => "true", + "max_results" => "10", + } search_post = Fabricate(:post, topic: topic_with_tags) + private_search_post = Fabricate(:post, topic: Fabricate(:topic, category: private_category)) + private_search_post.topic.tags = [tag_funny] + private_search_post.topic.save! _bot_post = Fabricate(:post) @@ -46,18 +62,28 @@ RSpec.describe DiscourseAi::AiBot::Tools::Search do bot_user: bot_user, llm: llm, context: { + user: user, }, ) + expect(search.options[:base_query]).to eq("#funny") + expect(search.options[:search_private]).to eq(true) + expect(search.options[:max_results]).to eq(10) + results = search.invoke(&progress_blk) expect(results[:rows].length).to eq(1) + GroupUser.create!(group: group, user: user) + + results = search.invoke(&progress_blk) + expect(results[:rows].length).to eq(2) + search_post.topic.tags = [] search_post.topic.save! - # no longer has the tag funny + # no longer has the tag funny, but secure one does results = search.invoke(&progress_blk) - expect(results[:rows].length).to eq(0) + expect(results[:rows].length).to eq(1) end it "can handle no results" do diff --git a/spec/requests/admin/ai_personas_controller_spec.rb b/spec/requests/admin/ai_personas_controller_spec.rb index e1f42622..565a44a2 100644 --- a/spec/requests/admin/ai_personas_controller_spec.rb +++ b/spec/requests/admin/ai_personas_controller_spec.rb @@ -75,6 +75,12 @@ RSpec.describe DiscourseAi::Admin::AiPersonasController do "description" => I18n.t("discourse_ai.ai_bot.command_options.search.max_results.description"), }, + "search_private" => { + "type" => "boolean", + "name" => I18n.t("discourse_ai.ai_bot.command_options.search.search_private.name"), + "description" => + I18n.t("discourse_ai.ai_bot.command_options.search.search_private.description"), + }, }, )