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.

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 = <template>
+export default class AiPersonaCommandOptionEditor extends Component {
-  <div class="control-group ai-persona-command-option-editor">
+  get isBoolean() {
-    <label>
+    return this.args.option.type === "boolean";
-      {{@option.name}}
+  }
-    </label>
-    <div class="">
-      <Input @value={{@option.value.value}} />
-    </div>
-    <div class="ai-persona-command-option-editor__instructions">
-      {{@option.description}}
-    </div>
-  </div>
-</template>;
 
-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";
+  }
+
+  <template>
+    <div class="control-group ai-persona-command-option-editor">
+      <label>
+        {{@option.name}}
+      </label>
+      <div class="">
+        {{#if this.isBoolean}}
+          <input
+            type="checkbox"
+            checked={{this.selectedValue}}
+            {{on "click" this.onCheckboxChange}}
+          />
+          {{@option.description}}
+        {{else}}
+          <Input @value={{@option.value.value}} />
+        {{/if}}
+      </div>
+      {{#unless this.isBoolean}}
+        <div class="ai-persona-command-option-editor__instructions">
+          {{@option.description}}
+        </div>
+      {{/unless}}
+    </div>
+  </template>
+}

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."

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.execute(safe_search_string, search_type: :full_page, guardian: guardian)
-              search_string.to_s + " status:public",
-              search_type: :full_page,
-              guardian: Guardian.new(),
-            )
 
           # 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

lib/ai_bot/tools/tool.rb

@@ -66,14 +66,20 @@ module DiscourseAi
         end
 
         def options
-          self
+          result = HashWithIndifferentAccess.new
-            .class
+          self.class.accepted_options.each do |option|
-            .accepted_options
+            val = @persona_options[option.name]
-            .reduce(HashWithIndifferentAccess.new) do |memo, option|
+            if val
-              val = @persona_options[option.name]
+              case option.type
-              memo[option.name] = val if val
+              when :boolean
-              memo
+                val = val == "true"
+              when :integer
+                val = val.to_i
+              end
+              result[option.name] = val
             end
+          end
+          result
         end
 
         def chain_next_response?

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

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"),
+          },
         },
       )