Allow AI tool to call search directly

Author: SamSaffron

lib/ai_bot/tool_runner.rb

@@ -73,6 +73,9 @@ def framework_script
         };
 
         const discourse = {
+          search: function(params) {
+            return _discourse_search(params);
+          },
           getPost: _discourse_get_post,
           getUser: _discourse_get_user,
           getPersona: function(name) {
@@ -341,6 +344,21 @@ def attach_discourse(mini_racer_context)
             end
           end,
         )
+
+        mini_racer_context.attach(
+          "_discourse_search",
+          ->(params) do
+            in_attached_function do
+              search_params = params.symbolize_keys
+              if search_params.delete(:with_private)
+                search_params[:current_user] = Discourse.system_user
+              end
+              search_params[:result_style] = :detailed
+              results = DiscourseAi::Utils::Search.perform_search(**search_params)
+              recursive_as_json(results)
+            end
+          end,
+        )
       end
 
       def attach_upload(mini_racer_context)

lib/ai_bot/tools/search.rb

@@ -34,7 +34,7 @@ def signature
                   enum: %w[latest latest_topic oldest views likes],
                 },
                 {
-                  name: "limit",
+                  name: "max_results",
                   description:
                     "limit number of results returned (generally prefer to just keep to default)",
                   type: "integer",
@@ -103,102 +103,38 @@ def search_query
 
         def invoke
           search_terms = []
-
           search_terms << options[:base_query] if options[:base_query].present?
-          search_terms << search_query.strip if search_query.present?
+          search_terms << search_query if search_query.present?
           search_args.each { |key, value| search_terms << "#{key}:#{value}" if value.present? }
 
-          guardian = nil
-          if options[:search_private] && context[:user]
-            guardian = Guardian.new(context[:user])
-          else
-            guardian = Guardian.new
-            search_terms << "status:public"
-          end
-
-          search_string = search_terms.join(" ").to_s
-          @last_query = search_string
-
-          yield(I18n.t("discourse_ai.ai_bot.searching", query: search_string))
+          @last_query = search_terms.join(" ").to_s
 
-          results = ::Search.execute(search_string, search_type: :full_page, guardian: guardian)
+          yield(I18n.t("discourse_ai.ai_bot.searching", query: @last_query))
 
           max_results = calculate_max_results(llm)
-          results_limit = parameters[:limit] || max_results
-          results_limit = max_results if parameters[:limit].to_i > max_results
-
-          should_try_semantic_search =
-            SiteSetting.ai_embeddings_semantic_search_enabled && search_query.present?
-
-          max_semantic_results = max_results / 4
-          results_limit = results_limit - max_semantic_results if should_try_semantic_search
-
-          posts = results&.posts || []
-          posts = posts[0..results_limit.to_i - 1]
-
-          if should_try_semantic_search
-            semantic_search = DiscourseAi::Embeddings::SemanticSearch.new(guardian)
-            topic_ids = Set.new(posts.map(&:topic_id))
-
-            search = ::Search.new(search_string, guardian: guardian)
-
-            results = nil
-            begin
-              results = semantic_search.search_for_topics(search.term)
-            rescue => e
-              Discourse.warn_exception(e, message: "Semantic search failed")
-            end
-
-            if results
-              results = search.apply_filters(results)
-
-              results.each do |post|
-                next if topic_ids.include?(post.topic_id)
-
-                topic_ids << post.topic_id
-                posts << post
-
-                break if posts.length >= max_results
-              end
-            end
+          if parameters[:max_results].to_i > 0
+            max_results = [parameters[:max_results].to_i, max_results].min
           end
 
-          @last_num_results = posts.length
-          # this is the general pattern from core
-          # if there are millions of hidden tags it may fail
-          hidden_tags = nil
-
-          if posts.blank?
-            { args: parameters, rows: [], instruction: "nothing was found, expand your search" }
-          else
-            format_results(posts, args: parameters) do |post|
-              category_names = [
-                post.topic.category&.parent_category&.name,
-                post.topic.category&.name,
-              ].compact.join(" > ")
-              row = {
-                title: post.topic.title,
-                url: Discourse.base_path + post.url,
-                username: post.user&.username,
-                excerpt: post.excerpt,
-                created: post.created_at,
-                category: category_names,
-                likes: post.like_count,
-                topic_views: post.topic.views,
-                topic_likes: post.topic.like_count,
-                topic_replies: post.topic.posts_count - 1,
-              }
-
-              if SiteSetting.tagging_enabled
-                hidden_tags ||= DiscourseTagging.hidden_tag_names
-                # using map over pluck to avoid n+1 (assuming caller preloading)
-                tags = post.topic.tags.map(&:name) - hidden_tags
-                row[:tags] = tags.join(", ") if tags.present?
-              end
-
-              row
-            end
-          end
+          search_query_with_base = [options[:base_query], search_query].compact.join(" ").strip
+
+          results =
+            DiscourseAi::Utils::Search.perform_search(
+              search_query: search_query_with_base,
+              category: parameters[:category],
+              user: parameters[:user],
+              order: parameters[:order],
+              max_posts: parameters[:max_posts],
+              tags: parameters[:tags],
+              before: parameters[:before],
+              after: parameters[:after],
+              status: parameters[:status],
+              max_results: max_results,
+              current_user: options[:search_private] ? context[:user] : nil,
+            )
+
+          @last_num_results = results[:rows]&.length || 0
+          results
         end
 
         protected

lib/utils/search.rb

@@ -0,0 +1,151 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module Utils
+    class Search
+      def self.perform_search(
+        search_query: nil,
+        category: nil,
+        user: nil,
+        order: nil,
+        max_posts: nil,
+        tags: nil,
+        before: nil,
+        after: nil,
+        status: nil,
+        hyde: true,
+        max_results: 20,
+        current_user: nil,
+        result_style: :compact
+      )
+        search_terms = []
+
+        search_terms << search_query.strip if search_query.present?
+        search_terms << "category:#{category}" if category.present?
+        search_terms << "user:#{user}" if user.present?
+        search_terms << "order:#{order}" if order.present?
+        search_terms << "max_posts:#{max_posts}" if max_posts.present?
+        search_terms << "tags:#{tags}" if tags.present?
+        search_terms << "before:#{before}" if before.present?
+        search_terms << "after:#{after}" if after.present?
+        search_terms << "status:#{status}" if status.present?
+
+        guardian = Guardian.new(current_user)
+
+        search_string = search_terms.join(" ").to_s
+
+        results = ::Search.execute(search_string, search_type: :full_page, guardian: guardian)
+        results_limit = max_results
+
+        should_try_semantic_search =
+          SiteSetting.ai_embeddings_semantic_search_enabled && search_query.present?
+
+        max_semantic_results = max_results / 4
+        results_limit = results_limit - max_semantic_results if should_try_semantic_search
+
+        posts = results&.posts || []
+        posts = posts[0..results_limit.to_i - 1]
+
+        if should_try_semantic_search
+          semantic_search = DiscourseAi::Embeddings::SemanticSearch.new(guardian)
+          topic_ids = Set.new(posts.map(&:topic_id))
+
+          search = ::Search.new(search_string, guardian: guardian)
+
+          semantic_results = nil
+          begin
+            semantic_results = semantic_search.search_for_topics(search.term, hyde: hyde)
+          rescue => e
+            Discourse.warn_exception(e, message: "Semantic search failed")
+          end
+
+          if semantic_results
+            semantic_results = search.apply_filters(semantic_results)
+
+            semantic_results.each do |post|
+              next if topic_ids.include?(post.topic_id)
+
+              topic_ids << post.topic_id
+              posts << post
+
+              break if posts.length >= max_results
+            end
+          end
+        end
+
+        hidden_tags = nil
+
+        # Construct search_args hash for consistent return format
+        search_args = {
+          search_query: search_query,
+          category: category,
+          user: user,
+          order: order,
+          max_posts: max_posts,
+          tags: tags,
+          before: before,
+          after: after,
+          status: status,
+          max_results: max_results,
+        }.compact
+
+        if posts.blank?
+          { args: search_args, rows: [], instruction: "nothing was found, expand your search" }
+        else
+          format_results(posts, args: search_args, result_style: result_style) do |post|
+            category_names = [
+              post.topic.category&.parent_category&.name,
+              post.topic.category&.name,
+            ].compact.join(" > ")
+            row = {
+              title: post.topic.title,
+              url: Discourse.base_path + post.url,
+              username: post.user&.username,
+              excerpt: post.excerpt,
+              created: post.created_at,
+              category: category_names,
+              likes: post.like_count,
+              topic_views: post.topic.views,
+              topic_likes: post.topic.like_count,
+              topic_replies: post.topic.posts_count - 1,
+            }
+
+            if SiteSetting.tagging_enabled
+              hidden_tags ||= DiscourseTagging.hidden_tag_names
+              tags = post.topic.tags.map(&:name) - hidden_tags
+              row[:tags] = tags.join(", ") if tags.present?
+            end
+
+            row
+          end
+        end
+      end
+
+      private
+
+      def self.format_results(rows, args: nil, result_style:)
+        rows = rows&.map { |row| yield row } if block_given?
+
+        if result_style == :compact
+          index = -1
+          column_indexes = {}
+
+          rows =
+            rows&.map do |data|
+              new_row = []
+              data.each do |key, value|
+                found_index = column_indexes[key.to_s] ||= (index += 1)
+                new_row[found_index] = value
+              end
+              new_row
+            end
+          column_names = column_indexes.keys
+        end
+
+        result = { column_names: column_names, rows: rows }
+        result[:args] = args if args
+        result
+      end
+    end
+  end
+end

spec/lib/modules/ai_bot/tools/search_spec.rb

@@ -98,7 +98,7 @@
 
       results = search.invoke(&progress_blk)
 
-      expect(results[:args]).to eq({ search_query: "ABDDCDCEDGDG", order: "fake" })
+      expect(results[:args]).to eq({ search_query: "ABDDCDCEDGDG", order: "fake", max_results: 60 })
       expect(results[:rows]).to eq([])
     end
 
@@ -131,7 +131,9 @@
             search.invoke(&progress_blk)
           end
 
-        expect(results[:args]).to eq({ search_query: "hello world, sam", status: "public" })
+        expect(results[:args]).to eq(
+          { max_results: 60, search_query: "hello world, sam", status: "public" },
+        )
         expect(results[:rows].length).to eq(1)
 
         # it also works with no query
@@ -174,6 +176,7 @@
               [param[:name], "test"]
             end
           end
+          .compact
           .to_h
           .symbolize_keys