From 944fd6569c80858eb251a0d8af79a65f3ecd5f31 Mon Sep 17 00:00:00 2001
From: Keegan George <kgeorge13@gmail.com>
Date: Thu, 1 Feb 2024 14:58:04 -0800
Subject: [PATCH] DEV: Add granular control for AI composer helper features
 (#458)

---
 .../ai-category-suggestion.gjs                |  2 +-
 .../ai-tag-suggestion.gjs                     |  2 +-
 .../ai-title-suggestion.gjs                   |  2 +-
 .../after-d-editor/ai-helper-context-menu.js  |  2 +-
 .../discourse/lib/show-ai-helper.js           |  8 ++++---
 config/locales/server.en.yml                  |  1 +
 config/settings.yml                           | 10 +++++++++
 .../ai_helper/ai_composer_helper_spec.rb      | 21 +++++++++++++++++++
 8 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/assets/javascripts/discourse/connectors/after-composer-category-input/ai-category-suggestion.gjs b/assets/javascripts/discourse/connectors/after-composer-category-input/ai-category-suggestion.gjs
index 2e3a0d27..50d6703e 100644
--- a/assets/javascripts/discourse/connectors/after-composer-category-input/ai-category-suggestion.gjs
+++ b/assets/javascripts/discourse/connectors/after-composer-category-input/ai-category-suggestion.gjs
@@ -5,7 +5,7 @@ import { showComposerAIHelper } from "../../lib/show-ai-helper";
 
 export default class AiCategorySuggestion extends Component {
   static shouldRender(outletArgs, helper) {
-    return showComposerAIHelper(outletArgs, helper);
+    return showComposerAIHelper(outletArgs, helper, "suggestions");
   }
 
   @service siteSettings;
diff --git a/assets/javascripts/discourse/connectors/after-composer-tag-input/ai-tag-suggestion.gjs b/assets/javascripts/discourse/connectors/after-composer-tag-input/ai-tag-suggestion.gjs
index 1c63a85e..d79697c4 100644
--- a/assets/javascripts/discourse/connectors/after-composer-tag-input/ai-tag-suggestion.gjs
+++ b/assets/javascripts/discourse/connectors/after-composer-tag-input/ai-tag-suggestion.gjs
@@ -5,7 +5,7 @@ import { showComposerAIHelper } from "../../lib/show-ai-helper";
 
 export default class AiTagSuggestion extends Component {
   static shouldRender(outletArgs, helper) {
-    return showComposerAIHelper(outletArgs, helper);
+    return showComposerAIHelper(outletArgs, helper, "suggestions");
   }
 
   @service siteSettings;
diff --git a/assets/javascripts/discourse/connectors/after-composer-title-input/ai-title-suggestion.gjs b/assets/javascripts/discourse/connectors/after-composer-title-input/ai-title-suggestion.gjs
index 9dc95850..69323c19 100644
--- a/assets/javascripts/discourse/connectors/after-composer-title-input/ai-title-suggestion.gjs
+++ b/assets/javascripts/discourse/connectors/after-composer-title-input/ai-title-suggestion.gjs
@@ -4,7 +4,7 @@ import { showComposerAIHelper } from "../../lib/show-ai-helper";
 
 export default class AiTitleSuggestion extends Component {
   static shouldRender(outletArgs, helper) {
-    return showComposerAIHelper(outletArgs, helper);
+    return showComposerAIHelper(outletArgs, helper, "suggestions");
   }
 
   <template>
diff --git a/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.js b/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.js
index 22dbc5cb..2b4a3fd9 100644
--- a/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.js
+++ b/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.js
@@ -12,7 +12,7 @@ import { showComposerAIHelper } from "../../lib/show-ai-helper";
 
 export default class AiHelperContextMenu extends Component {
   static shouldRender(outletArgs, helper) {
-    return showComposerAIHelper(outletArgs, helper);
+    return showComposerAIHelper(outletArgs, helper, "context_menu");
   }
 
   @service currentUser;
diff --git a/assets/javascripts/discourse/lib/show-ai-helper.js b/assets/javascripts/discourse/lib/show-ai-helper.js
index 84edb754..13e16076 100644
--- a/assets/javascripts/discourse/lib/show-ai-helper.js
+++ b/assets/javascripts/discourse/lib/show-ai-helper.js
@@ -1,16 +1,18 @@
-export function showComposerAIHelper(outletArgs, helper) {
+export function showComposerAIHelper(outletArgs, helper, featureType) {
   const enableHelper = _helperEnabled(helper.siteSettings);
   const enableAssistant = _canUseAssistant(
     helper.currentUser,
     _findAllowedGroups(helper.siteSettings.ai_helper_allowed_groups)
   );
   const canShowInPM = helper.siteSettings.ai_helper_allowed_in_pm;
+  const enableFeature =
+    helper.siteSettings.ai_helper_enabled_features.includes(featureType);
 
   if (outletArgs?.composer?.privateMessage) {
-    return enableHelper && enableAssistant && canShowInPM;
+    return enableHelper && enableAssistant && canShowInPM && enableFeature;
   }
 
-  return enableHelper && enableAssistant;
+  return enableHelper && enableAssistant && enableFeature;
 }
 
 export function showPostAIHelper(outletArgs, helper) {
diff --git a/config/locales/server.en.yml b/config/locales/server.en.yml
index 023ab611..c5886d88 100644
--- a/config/locales/server.en.yml
+++ b/config/locales/server.en.yml
@@ -60,6 +60,7 @@ en:
     ai_helper_automatic_chat_thread_title_delay: "Delay in minutes before the AI helper automatically sets the chat thread title."
     ai_helper_automatic_chat_thread_title: "Automatically set the chat thread titles based on thread contents."
     ai_helper_illustrate_post_model: "Model to use for the composer AI helper's illustrate post feature"
+    ai_helper_enabled_features: "Select the features to enable in the AI helper."
 
     ai_embeddings_enabled: "Enable the embeddings module."
     ai_embeddings_discourse_service_api_endpoint: "URL where the API is running for the embeddings module"
diff --git a/config/settings.yml b/config/settings.yml
index fe8be23b..64585591 100644
--- a/config/settings.yml
+++ b/config/settings.yml
@@ -212,6 +212,16 @@ discourse_ai:
       - stable_diffusion_xl
       - dall_e_3
       - disabled
+  ai_helper_enabled_features:
+    client: true
+    default: "suggestions|context_menu"
+    type: list
+    list_type: compact
+    allow_any: false
+    refresh: true
+    choices:
+      - "suggestions"
+      - "context_menu"
 
   ai_embeddings_enabled:
     default: false
diff --git a/spec/system/ai_helper/ai_composer_helper_spec.rb b/spec/system/ai_helper/ai_composer_helper_spec.rb
index dff111f6..d0edaf48 100644
--- a/spec/system/ai_helper/ai_composer_helper_spec.rb
+++ b/spec/system/ai_helper/ai_composer_helper_spec.rb
@@ -412,4 +412,25 @@ RSpec.describe "AI Composer helper", type: :system, js: true do
       expect(ai_suggestion_dropdown).to have_no_suggestion_button
     end
   end
+
+  context "when suggestion features are disabled" do
+    let(:mode) { CompletionPrompt::GENERATE_TITLES }
+    before { SiteSetting.ai_helper_enabled_features = "context_menu" }
+
+    it "does not show suggestion buttons in the composer" do
+      visit("/latest")
+      page.find("#create-topic").click
+      composer.fill_content(input)
+      expect(ai_suggestion_dropdown).to have_no_suggestion_button
+    end
+  end
+
+  context "when context menu feature is disabled" do
+    before { SiteSetting.ai_helper_enabled_features = "suggestions" }
+
+    it "does not show context menu in the composer" do
+      trigger_context_menu(input)
+      expect(ai_helper_context_menu).to have_no_context_menu
+    end
+  end
 end