From f1283e156d48b7155d2025061e70f6b5df32c662 Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Wed, 23 Oct 2024 16:55:10 +1100
Subject: [PATCH] FEATURE: allow scoping of google tool queries (#852)

This allows to simply scope search results to specific domains and prepend arbitrary snippets to searches made
---
 config/locales/server.en.yml                 |  4 +++
 lib/ai_bot/tools/google.rb                   |  9 +++++
 spec/lib/modules/ai_bot/tools/google_spec.rb | 35 ++++++++++++++++----
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/config/locales/server.en.yml b/config/locales/server.en.yml
index 538da773..2c2033a0 100644
--- a/config/locales/server.en.yml
+++ b/config/locales/server.en.yml
@@ -205,6 +205,10 @@ en:
       summarizing: "Summarizing topic"
       searching: "Searching for: '%{query}'"
       tool_options:
+        google:
+          base_query:
+            name: "Base Search Query"
+            description: "Base query to use when searching. Examples: 'site:example.com' will only include results from example.com, before:2022-01-01 will only includes results from 2021 and earlier. This text is prepended to the search query."
         read:
           read_private:
             name: "Read Private"
diff --git a/lib/ai_bot/tools/google.rb b/lib/ai_bot/tools/google.rb
index ba3ecb5d..77158882 100644
--- a/lib/ai_bot/tools/google.rb
+++ b/lib/ai_bot/tools/google.rb
@@ -23,15 +23,24 @@ module DiscourseAi
           "google"
         end
 
+        def self.accepted_options
+          [option(:base_query, type: :string)]
+        end
+
         def query
           parameters[:query].to_s.strip
         end
 
         def invoke
+          query = self.query
+
           yield(query)
 
           api_key = SiteSetting.ai_google_custom_search_api_key
           cx = SiteSetting.ai_google_custom_search_cx
+
+          query = "#{options[:base_query]} #{query}" if options[:base_query].present?
+
           escaped_query = CGI.escape(query)
           uri =
             URI(
diff --git a/spec/lib/modules/ai_bot/tools/google_spec.rb b/spec/lib/modules/ai_bot/tools/google_spec.rb
index 9f522cef..c39f94b0 100644
--- a/spec/lib/modules/ai_bot/tools/google_spec.rb
+++ b/spec/lib/modules/ai_bot/tools/google_spec.rb
@@ -7,13 +7,14 @@ RSpec.describe DiscourseAi::AiBot::Tools::Google do
   let(:progress_blk) { Proc.new {} }
   let(:search) { described_class.new({ query: "some search term" }, bot_user: bot_user, llm: llm) }
 
-  before { SiteSetting.ai_bot_enabled = true }
+  before do
+    SiteSetting.ai_bot_enabled = true
+    SiteSetting.ai_google_custom_search_api_key = "abc"
+    SiteSetting.ai_google_custom_search_cx = "cx"
+  end
 
   describe "#process" do
     it "will not explode if there are no results" do
-      SiteSetting.ai_google_custom_search_api_key = "abc"
-      SiteSetting.ai_google_custom_search_cx = "cx"
-
       json_text = { searchInformation: { totalResults: "0" } }.to_json
 
       stub_request(
@@ -27,10 +28,30 @@ RSpec.describe DiscourseAi::AiBot::Tools::Google do
       expect(info).to_not include("oops")
     end
 
-    it "can generate correct info" do
-      SiteSetting.ai_google_custom_search_api_key = "abc"
-      SiteSetting.ai_google_custom_search_cx = "cx"
+    it "supports base_query" do
+      base_query = "site:discourse.org"
 
+      search =
+        described_class.new(
+          { query: "some search term" },
+          bot_user: bot_user,
+          llm: llm,
+          persona_options: {
+            "base_query" => base_query,
+          },
+        )
+
+      json_text = { searchInformation: { totalResults: "0" } }.to_json
+
+      stub_request(
+        :get,
+        "https://www.googleapis.com/customsearch/v1?cx=cx&key=abc&num=10&q=site:discourse.org%20some%20search%20term",
+      ).to_return(status: 200, body: json_text, headers: {})
+
+      search.invoke(&progress_blk)
+    end
+
+    it "can generate correct info" do
       json_text = {
         searchInformation: {
           totalResults: "2",