From 74a7ac4a3d04726c34ab0b6fd9cfbf1b44a793f6 Mon Sep 17 00:00:00 2001
From: Keegan George <kgeorge13@gmail.com>
Date: Thu, 14 Dec 2023 08:47:20 -0800
Subject: [PATCH] FEATURE: Add custom prompts to post helper options (#355)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* FEATURE: Add custom prompts to post helper options

* 💄Make pretty

* 💄Make pretty!
---
 .../components/ai-helper-custom-prompt.gjs    | 46 ++++++++++
 .../after-d-editor/ai-helper-context-menu.hbs | 26 ++----
 .../ai-helper-options-menu.gjs                | 83 +++++++++++++------
 .../modules/ai-helper/common/ai-helper.scss   | 49 +++++------
 lib/ai_helper/assistant.rb                    |  4 +-
 5 files changed, 131 insertions(+), 77 deletions(-)
 create mode 100644 assets/javascripts/discourse/components/ai-helper-custom-prompt.gjs

diff --git a/assets/javascripts/discourse/components/ai-helper-custom-prompt.gjs b/assets/javascripts/discourse/components/ai-helper-custom-prompt.gjs
new file mode 100644
index 00000000..b94cf0bd
--- /dev/null
+++ b/assets/javascripts/discourse/components/ai-helper-custom-prompt.gjs
@@ -0,0 +1,46 @@
+import Component from "@glimmer/component";
+import { tracked } from "@glimmer/tracking";
+import { Input } from "@ember/component";
+import { action } from "@ember/object";
+import didInsert from "@ember/render-modifiers/modifiers/did-insert";
+import DButton from "discourse/components/d-button";
+import i18n from "discourse-common/helpers/i18n";
+import not from "truth-helpers/helpers/not";
+
+export default class AiHelperCustomPrompt extends Component {
+  <template>
+    <div class="ai-custom-prompt" {{didInsert this.setupCustomPrompt}}>
+      <Input
+        @value={{@value}}
+        placeholder={{i18n
+          "discourse_ai.ai_helper.context_menu.custom_prompt.placeholder"
+        }}
+        class="ai-custom-prompt__input"
+        @enter={{this.sendInput}}
+      />
+
+      <DButton
+        @class="ai-custom-prompt__submit btn-primary"
+        @icon="discourse-sparkles"
+        @action={{@submit}}
+        @actionParam={{@promptArgs}}
+        @disabled={{not @value.length}}
+      />
+    </div>
+  </template>
+
+  @tracked _customPromptInput;
+
+  @action
+  setupCustomPrompt() {
+    this._customPromptInput = document.querySelector(
+      ".ai-custom-prompt__input"
+    );
+    this._customPromptInput.focus();
+  }
+
+  @action
+  sendInput() {
+    return this.args.submit(this.args.promptArgs);
+  }
+}
diff --git a/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.hbs b/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.hbs
index b63cad7a..14843893 100644
--- a/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.hbs
+++ b/assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.hbs
@@ -17,27 +17,11 @@
         <ul class="ai-helper-context-menu__options">
           {{#each this.helperOptions as |option|}}
             {{#if (eq option.name "custom_prompt")}}
-              <div
-                class="ai-custom-prompt"
-                {{did-insert this.setupCustomPrompt}}
-              >
-                <Input
-                  @value={{this.customPromptValue}}
-                  placeholder={{i18n
-                    "discourse_ai.ai_helper.context_menu.custom_prompt.placeholder"
-                  }}
-                  class="ai-custom-prompt__input"
-                  @enter={{action (fn this.updateSelected option)}}
-                />
-
-                <DButton
-                  @class="ai-custom-prompt__submit btn-primary"
-                  @icon="discourse-sparkles"
-                  @action={{this.updateSelected}}
-                  @actionParam={{option}}
-                  @disabled={{not this.customPromptValue.length}}
-                />
-              </div>
+              <AiHelperCustomPrompt
+                @value={{this.customPromptValue}}
+                @promptArgs={{option}}
+                @submit={{this.updateSelected}}
+              />
             {{else}}
               <li data-name={{option.translated_name}} data-value={{option.id}}>
                 <DButton
diff --git a/assets/javascripts/discourse/connectors/post-text-buttons/ai-helper-options-menu.gjs b/assets/javascripts/discourse/connectors/post-text-buttons/ai-helper-options-menu.gjs
index 2062eed3..a6b106e5 100644
--- a/assets/javascripts/discourse/connectors/post-text-buttons/ai-helper-options-menu.gjs
+++ b/assets/javascripts/discourse/connectors/post-text-buttons/ai-helper-options-menu.gjs
@@ -11,6 +11,7 @@ import { cook } from "discourse/lib/text";
 import { bind } from "discourse-common/utils/decorators";
 import eq from "truth-helpers/helpers/eq";
 import not from "truth-helpers/helpers/not";
+import AiHelperCustomPrompt from "../../components/ai-helper-custom-prompt";
 import AiHelperLoading from "../../components/ai-helper-loading";
 import { showPostAIHelper } from "../../lib/show-ai-helper";
 
@@ -19,12 +20,14 @@ export default class AIHelperOptionsMenu extends Component {
     return showPostAIHelper(outletArgs, helper);
   }
   @service messageBus;
+  @service siteSettings;
+  @service currentUser;
   @tracked helperOptions = [];
   @tracked menuState = this.MENU_STATES.triggers;
   @tracked loading = false;
   @tracked suggestion = "";
   @tracked showMainButtons = true;
-
+  @tracked customPromptValue = "";
   @tracked copyButtonIcon = "copy";
   @tracked copyButtonLabel = "discourse_ai.ai_helper.post_options_menu.copy";
 
@@ -96,18 +99,23 @@ export default class AIHelperOptionsMenu extends Component {
       this._activeAIRequest = ajax("/discourse-ai/ai-helper/suggest", {
         method: "POST",
         data: {
-          mode: option.value,
+          mode: option.id,
           text: this.args.outletArgs.data.quoteState.buffer,
-          custom_prompt: "",
+          custom_prompt: this.customPromptValue,
         },
       });
     }
 
     if (option.name !== "Explain") {
-      this._activeAIRequest.catch(popupAjaxError).finally(() => {
-        this.loading = false;
-        this.menuState = this.MENU_STATES.result;
-      });
+      this._activeAIRequest
+        .then(({ suggestions }) => {
+          this.suggestion = suggestions[0];
+        })
+        .catch(popupAjaxError)
+        .finally(() => {
+          this.loading = false;
+          this.menuState = this.MENU_STATES.result;
+        });
     }
 
     return this._activeAIRequest;
@@ -142,15 +150,32 @@ export default class AIHelperOptionsMenu extends Component {
   async loadPrompts() {
     let prompts = await ajax("/discourse-ai/ai-helper/prompts");
 
-    this.helperOptions = prompts
-      .filter((item) => item.location.includes("post"))
-      .map((p) => {
-        return {
-          name: p.translated_name,
-          value: p.id,
-          icon: p.icon,
-        };
-      });
+    prompts = prompts.filter((item) => item.location.includes("post"));
+
+    // Find the custom_prompt object and move it to the beginning of the array
+    const customPromptIndex = prompts.findIndex(
+      (p) => p.name === "custom_prompt"
+    );
+
+    if (customPromptIndex !== -1) {
+      const customPrompt = prompts.splice(customPromptIndex, 1)[0];
+      prompts.unshift(customPrompt);
+    }
+
+    if (!this._showUserCustomPrompts()) {
+      prompts = prompts.filter((p) => p.name !== "custom_prompt");
+    }
+
+    this.helperOptions = prompts;
+  }
+
+  _showUserCustomPrompts() {
+    const allowedGroups =
+      this.siteSettings?.ai_helper_custom_prompts_allowed_groups
+        .split("|")
+        .map((id) => parseInt(id, 10));
+
+    return this.currentUser?.groups.some((g) => allowedGroups.includes(g.id));
   }
 
   <template>
@@ -170,15 +195,23 @@ export default class AIHelperOptionsMenu extends Component {
       {{else if (eq this.menuState this.MENU_STATES.options)}}
         <div class="ai-post-helper__options">
           {{#each this.helperOptions as |option|}}
-            <DButton
-              @class="btn-flat"
-              @icon={{option.icon}}
-              @translatedLabel={{option.name}}
-              @action={{this.performAISuggestion}}
-              @actionParam={{option}}
-              data-name={{option.name}}
-              data-value={{option.value}}
-            />
+            {{#if (eq option.name "custom_prompt")}}
+              <AiHelperCustomPrompt
+                @value={{this.customPromptValue}}
+                @promptArgs={{option}}
+                @submit={{this.performAISuggestion}}
+              />
+            {{else}}
+              <DButton
+                @class="btn-flat"
+                @icon={{option.icon}}
+                @translatedLabel={{option.translated_name}}
+                @action={{this.performAISuggestion}}
+                @actionParam={{option}}
+                data-name={{option.name}}
+                data-value={{option.value}}
+              />
+            {{/if}}
           {{/each}}
         </div>
 
diff --git a/assets/stylesheets/modules/ai-helper/common/ai-helper.scss b/assets/stylesheets/modules/ai-helper/common/ai-helper.scss
index 530c8cf0..4361d8a5 100644
--- a/assets/stylesheets/modules/ai-helper/common/ai-helper.scss
+++ b/assets/stylesheets/modules/ai-helper/common/ai-helper.scss
@@ -108,39 +108,19 @@
     align-items: center;
     flex-flow: row wrap;
   }
+}
 
-  &__custom-prompt {
-    display: flex;
-    flex-flow: row wrap;
-    padding: 0.5rem;
+.ai-custom-prompt {
+  display: flex;
+  gap: 0.25rem;
+  margin-bottom: 0.5rem;
 
-    &-header {
-      margin-bottom: 0.5rem;
+  &__input[type="text"] {
+    border-color: var(--primary-400);
+    margin-bottom: 0;
 
-      .btn {
-        padding: 0;
-      }
-    }
-
-    .ai-custom-prompt-input {
-      min-height: 90px;
-      width: 100%;
-    }
-  }
-
-  .ai-custom-prompt {
-    display: flex;
-    gap: 0.25rem;
-    margin-bottom: 0.5rem;
-
-    &__input {
-      background: var(--primary-low);
-      border-color: var(--primary-low);
-      margin-bottom: 0;
-
-      &::placeholder {
-        color: var(--primary-medium);
-      }
+    &::placeholder {
+      color: var(--primary-medium);
     }
   }
 }
@@ -326,6 +306,15 @@
     align-items: flex-start;
     gap: 0.25rem;
     justify-content: flex-start;
+
+    > button:last-child {
+      margin-bottom: 0.5rem;
+    }
+
+    .ai-custom-prompt {
+      padding: 0.5rem;
+      margin-bottom: 0;
+    }
   }
 
   &__suggestion {
diff --git a/lib/ai_helper/assistant.rb b/lib/ai_helper/assistant.rb
index 139ceacb..573fec35 100644
--- a/lib/ai_helper/assistant.rb
+++ b/lib/ai_helper/assistant.rb
@@ -85,6 +85,8 @@ module DiscourseAi
           </input>
           <output>
           </output>
+          <result>
+          </result>
         ]
 
         result.dup.tap { |dup_result| tags_to_remove.each { |tag| dup_result.gsub!(tag, "") } }
@@ -130,7 +132,7 @@ module DiscourseAi
         when "tone"
           %w[composer]
         when "custom_prompt"
-          %w[composer]
+          %w[composer post]
         when "rewrite"
           %w[composer]
         when "explain"