FEATURE: Add streaming to post AI helper's explain option (#344)

Co-authored-by: Rafael dos Santos Silva <xfalcox@gmail.com>
Co-authored-by: Roman Rizzi <roman@discourse.org>

app/controllers/discourse_ai/ai_helper/assistant_controller.rb

@@ -101,12 +101,14 @@ module DiscourseAi
 
         raise Discourse::InvalidParameters.new(:post_id) unless post
 
-        render json:
+        Jobs.enqueue(
-                 DiscourseAi::AiHelper::TopicHelper.new(current_user).explain(
+          :stream_post_helper,
-                   term_to_explain,
+          post_id: post.id,
-                   post,
+          user_id: current_user.id,
-                 ),
+          term_to_explain: term_to_explain,
-               status: 200
+        )
+
+        render json: { success: true }, status: 200
       rescue DiscourseAi::Completions::Endpoints::Base::CompletionFailed => e
         render_json_error I18n.t("discourse_ai.ai_helper.errors.completion_request_failed"),
                           status: 502

app/jobs/regular/stream_post_helper.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+module Jobs
+  class StreamPostHelper < ::Jobs::Base
+    sidekiq_options retry: false
+
+    def execute(args)
+      return unless post = Post.includes(:topic).find_by(id: args[:post_id])
+      return unless user = User.find_by(id: args[:user_id])
+      return unless args[:term_to_explain]
+
+      topic = post.topic
+      reply_to = post.reply_to_post
+
+      guardian = Guardian.new(user)
+      return unless guardian.can_see?(post)
+
+      prompt = CompletionPrompt.enabled_by_name("explain")
+
+      input = <<~TEXT
+      <term>#{args[:term_to_explain]}</term>
+      <context>#{post.raw}</context>
+          <topic>#{topic.title}</topic>
+          #{reply_to ? "<replyTo>#{reply_to.raw}</replyTo>" : nil}
+        TEXT
+
+      DiscourseAi::AiHelper::Assistant.new.stream_prompt(
+        prompt,
+        input,
+        user,
+        "/discourse-ai/ai-helper/explain/#{post.id}",
+      )
+    end
+  end
+end

assets/javascripts/discourse/components/ai-helper-loading.gjs

@@ -0,0 +1,19 @@
+import DButton from "discourse/components/d-button";
+import i18n from "discourse-common/helpers/i18n";
+
+const AiHelperLoading = <template>
+  <div class="ai-helper-context-menu__loading">
+    <div class="dot-falling"></div>
+    <span>
+      {{i18n "discourse_ai.ai_helper.context_menu.loading"}}
+    </span>
+    <DButton
+      @icon="times"
+      @title="discourse_ai.ai_helper.context_menu.cancel"
+      @action={{@cancel}}
+      class="btn-flat cancel-request"
+    />
+  </div>
+</template>;
+
+export default AiHelperLoading;

assets/javascripts/discourse/connectors/after-d-editor/ai-helper-context-menu.hbs

@@ -53,18 +53,7 @@
         </ul>
 
       {{else if (eq this.menuState this.CONTEXT_MENU_STATES.loading)}}
-        <div class="ai-helper-context-menu__loading">
+        <AiHelperLoading @cancel={{this.cancelAIAction}} />
-          <div class="dot-falling"></div>
-          <span>
-            {{i18n "discourse_ai.ai_helper.context_menu.loading"}}
-          </span>
-          <DButton
-            @icon="times"
-            @title="discourse_ai.ai_helper.context_menu.cancel"
-            @action={{this.cancelAIAction}}
-            class="btn-flat cancel-request"
-          />
-        </div>
 
       {{else if (eq this.menuState this.CONTEXT_MENU_STATES.review)}}
         <ul class="ai-helper-context-menu__review">

assets/javascripts/discourse/connectors/post-text-buttons/ai-helper-options-menu.gjs

@@ -1,20 +1,23 @@
 import Component from "@glimmer/component";
 import { tracked } from "@glimmer/tracking";
 import { action } from "@ember/object";
+import didInsert from "@ember/render-modifiers/modifiers/did-insert";
+import willDestroy from "@ember/render-modifiers/modifiers/will-destroy";
+import { inject as service } from "@ember/service";
 import DButton from "discourse/components/d-button";
 import { ajax } from "discourse/lib/ajax";
 import { popupAjaxError } from "discourse/lib/ajax-error";
-import I18n from "I18n";
+import { bind } from "discourse-common/utils/decorators";
 import eq from "truth-helpers/helpers/eq";
 import not from "truth-helpers/helpers/not";
+import AiHelperLoading from "../../components/ai-helper-loading";
 import { showPostAIHelper } from "../../lib/show-ai-helper";
 
-const i18n = I18n.t.bind(I18n);
-
 export default class AIHelperOptionsMenu extends Component {
   static shouldRender(outletArgs, helper) {
     return showPostAIHelper(outletArgs, helper);
   }
+  @service messageBus;
   @tracked helperOptions = [];
   @tracked menuState = this.MENU_STATES.triggers;
   @tracked loading = false;
@@ -47,12 +50,38 @@ export default class AIHelperOptionsMenu extends Component {
     this.menuState = this.MENU_STATES.options;
   }
 
+  @bind
+  subscribe() {
+    const channel = `/discourse-ai/ai-helper/explain/${this.args.outletArgs.data.quoteState.postId}`;
+    this.messageBus.subscribe(channel, this._updateResult);
+  }
+
+  @bind
+  unsubscribe() {
+    this.messageBus.unsubscribe(
+      "/discourse-ai/ai-helper/explain/*",
+      this._updateResult
+    );
+  }
+
+  @bind
+  _updateResult(result) {
+    const suggestion = result.result;
+
+    if (suggestion.length > 0) {
+      this.suggestion = suggestion;
+    }
+  }
+
   @action
   async performAISuggestion(option) {
     this.menuState = this.MENU_STATES.loading;
 
     if (option.name === "Explain") {
-      this._activeAIRequest = ajax("/discourse-ai/ai-helper/explain", {
+      this.menuState = this.MENU_STATES.result;
+
+      const fetchUrl = `/discourse-ai/ai-helper/explain`;
+      this._activeAIRequest = ajax(fetchUrl, {
         method: "POST",
         data: {
           mode: option.value,
@@ -71,6 +100,7 @@ export default class AIHelperOptionsMenu extends Component {
       });
     }
 
+    if (option.name !== "Explain") {
       this._activeAIRequest
         .then(({ suggestions }) => {
           this.suggestion = suggestions[0];
@@ -80,6 +110,7 @@ export default class AIHelperOptionsMenu extends Component {
           this.loading = false;
           this.menuState = this.MENU_STATES.result;
         });
+    }
 
     return this._activeAIRequest;
   }
@@ -154,23 +185,24 @@ export default class AIHelperOptionsMenu extends Component {
         </div>
 
       {{else if (eq this.menuState this.MENU_STATES.loading)}}
-        <div class="ai-helper-context-menu__loading">
+        <AiHelperLoading @cancel={{this.cancelAIAction}} />
-          <div class="dot-falling"></div>
-          <span>
-            {{i18n "discourse_ai.ai_helper.context_menu.loading"}}
-          </span>
-          <DButton
-            @icon="times"
-            @title="discourse_ai.ai_helper.context_menu.cancel"
-            @action={{this.cancelAIAction}}
-            class="btn-flat cancel-request"
-          />
-        </div>
       {{else if (eq this.menuState this.MENU_STATES.result)}}
-        <div class="ai-post-helper__suggestion">
+        <div
+          class="ai-post-helper__suggestion"
+          {{didInsert this.subscribe}}
+          {{willDestroy this.unsubscribe}}
+        >
+          {{#if this.suggestion}}
             <div class="ai-post-helper__suggestion__text">
               {{this.suggestion}}
             </div>
+            <di class="ai-post-helper__suggestion__buttons">
+              <DButton
+                @class="btn-flat ai-post-helper__suggestion__cancel"
+                @icon="times"
+                @label="discourse_ai.ai_helper.post_options_menu.cancel"
+                @action={{this.cancelAIAction}}
+              />
               <DButton
                 @class="btn-flat ai-post-helper__suggestion__copy"
                 @icon={{this.copyButtonIcon}}
@@ -178,6 +210,10 @@ export default class AIHelperOptionsMenu extends Component {
                 @action={{this.copySuggestion}}
                 @disabled={{not this.suggestion}}
               />
+            </di>
+          {{else}}
+            <AiHelperLoading @cancel={{this.cancelAIAction}} />
+          {{/if}}
         </div>
       {{/if}}
     </div>

assets/stylesheets/modules/ai-helper/common/ai-helper.scss

@@ -334,8 +334,6 @@
     flex-direction: column;
 
     &__copy {
-      margin-top: 0.5rem;
-
       .d-icon-check {
         color: var(--success);
       }
@@ -344,5 +342,16 @@
     &__text {
       padding: 0.5rem;
     }
+
+    &__buttons {
+      display: flex;
+      align-items: center;
+      justify-content: stretch;
+      margin-top: 0.5rem;
+      gap: 0.5rem;
+      .btn {
+        width: 100%;
+      }
+    }
   }
 }

config/locales/client.en.yml

@@ -103,6 +103,7 @@ en:
           close: "Close"
           copy: "Copy"
           copied: "Copied!"
+          cancel: "Cancel"
       reviewables:
         model_used: "Model used:"
         accuracy: "Accuracy:"

lib/ai_helper/assistant.rb

@@ -24,29 +24,76 @@ module DiscourseAi
         end
       end
 
-      def generate_and_send_prompt(completion_prompt, input, user)
+      def generate_prompt(completion_prompt, input, user, &block)
         llm = DiscourseAi::Completions::Llm.proxy(SiteSetting.ai_helper_model)
-
         generic_prompt = completion_prompt.messages_with_input(input)
 
-        completion_result = llm.completion!(generic_prompt, user)
+        llm.completion!(generic_prompt, user, &block)
+      end
+
+      def generate_and_send_prompt(completion_prompt, input, user)
+        completion_result = generate_prompt(completion_prompt, input, user)
         result = { type: completion_prompt.prompt_type }
 
         result[:diff] = parse_diff(input, completion_result) if completion_prompt.diff?
 
         result[:suggestions] = (
           if completion_prompt.list?
-            parse_list(completion_result)
+            parse_list(completion_result).map { |suggestion| sanitize_result(suggestion) }
           else
-            [completion_result]
+            [sanitize_result(completion_result)]
           end
         )
 
         result
       end
 
+      def stream_prompt(completion_prompt, input, user, channel)
+        streamed_result = +""
+        start = Time.now
+
+        generate_prompt(completion_prompt, input, user) do |partial_response, cancel_function|
+          streamed_result << partial_response
+
+          # Throttle the updates
+          if (Time.now - start > 0.5) || Rails.env.test?
+            payload = { result: sanitize_result(streamed_result), done: false }
+            publish_update(channel, payload, user)
+            start = Time.now
+          end
+        end
+
+        sanitized_result = sanitize_result(streamed_result)
+        if sanitized_result.present?
+          publish_update(channel, { result: sanitized_result, done: true }, user)
+        end
+      end
+
       private
 
+      def sanitize_result(result)
+        tags_to_remove = %w[
+          <term>
+          </term>
+          <context>
+          </context>
+          <topic>
+          </topic>
+          <replyTo>
+          </replyTo>
+          <input>
+          </input>
+          <output>
+          </output>
+        ]
+
+        result.dup.tap { |dup_result| tags_to_remove.each { |tag| dup_result.gsub!(tag, "") } }
+      end
+
+      def publish_update(channel, payload, user)
+        MessageBus.publish(channel, payload, user_ids: [user.id])
+      end
+
       def icon_map(name)
         case name
         when "translate"

lib/ai_helper/topic_helper.rb

@@ -1,35 +0,0 @@
-# frozen_string_literal: true
-
-module DiscourseAi
-  module AiHelper
-    class TopicHelper
-      def initialize(user)
-        @user = user
-      end
-
-      def explain(term_to_explain, post)
-        return nil unless term_to_explain
-        return nil unless post
-
-        reply_to = post.reply_to_post
-        topic = post.topic
-
-        prompt = CompletionPrompt.enabled_by_name("explain")
-        raise Discourse::InvalidParameters.new(:mode) if !prompt
-
-        input = <<~TEXT
-          <term>#{term_to_explain}</term>
-          <context>#{post.raw}</context>
-          <topic>#{topic.title}</topic>
-          #{reply_to ? "<replyTo>#{reply_to.raw}</replyTo>" : nil}
-        TEXT
-
-        DiscourseAi::AiHelper::Assistant.new.generate_and_send_prompt(prompt, input, user)
-      end
-
-      private
-
-      attr_reader :user
-    end
-  end
-end

spec/jobs/regular/stream_post_helper.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+RSpec.describe Jobs::StreamPostHelper do
+  subject(:job) { described_class.new }
+
+  describe "#execute" do
+    fab!(:topic) { Fabricate(:topic) }
+    fab!(:post) do
+      Fabricate(
+        :post,
+        topic: topic,
+        raw:
+          "I like to eat pie. It is a very good dessert. Some people are wasteful by throwing pie at others but I do not do that. I always eat the pie.",
+      )
+    end
+    fab!(:user) { Fabricate(:leader) }
+
+    before do
+      Group.find(Group::AUTO_GROUPS[:trust_level_3]).add(user)
+      SiteSetting.composer_ai_helper_enabled = true
+    end
+
+    describe "validates params" do
+      it "does nothing if there is no post" do
+        messages =
+          MessageBus.track_publish("/discourse-ai/ai-helper/explain/#{post.id}") do
+            job.execute(post_id: nil, user_id: user.id, term_to_explain: "pie")
+          end
+
+        expect(messages).to be_empty
+      end
+
+      it "does nothing if there is no user" do
+        messages =
+          MessageBus.track_publish("/discourse-ai/ai-helper/explain/#{post.id}") do
+            job.execute(post_id: post.id, user_id: nil, term_to_explain: "pie")
+          end
+
+        expect(messages).to be_empty
+      end
+
+      it "does nothing if there is no term to explain" do
+        messages =
+          MessageBus.track_publish("/discourse-ai/ai-helper/explain/#{post.id}") do
+            job.execute(post_id: post.id, user_id: user.id, term_to_explain: nil)
+          end
+
+        expect(messages).to be_empty
+      end
+    end
+
+    it "publishes updates with a partial result" do
+      explanation =
+        "In this context, \"pie\" refers to a baked dessert typically consisting of a pastry crust and filling."
+
+      DiscourseAi::Completions::Llm.with_prepared_responses([explanation]) do
+        messages =
+          MessageBus.track_publish("/discourse-ai/ai-helper/explain/#{post.id}") do
+            job.execute(post_id: post.id, user_id: user.id, term_to_explain: "pie")
+          end
+
+        partial_result_update = messages.first.data
+        expect(partial_result_update[:done]).to eq(false)
+        expect(partial_result_update[:result]).to eq(explanation)
+      end
+    end
+
+    it "publishes a final update to signal we're donea" do
+      explanation =
+        "In this context, \"pie\" refers to a baked dessert typically consisting of a pastry crust and filling."
+
+      DiscourseAi::Completions::Llm.with_prepared_responses([explanation]) do
+        messages =
+          MessageBus.track_publish("/discourse-ai/ai-helper/explain/#{post.id}") do
+            job.execute(post_id: post.id, user_id: user.id, term_to_explain: "pie")
+          end
+
+        final_update = messages.last.data
+        expect(final_update[:done]).to eq(true)
+      end
+    end
+  end
+end

spec/system/ai_helper/ai_post_helper_spec.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-RSpec.describe "AI Composer helper", type: :system, js: true do
+RSpec.describe "AI Post helper", type: :system, js: true do
   fab!(:user) { Fabricate(:admin) }
   fab!(:non_member_group) { Fabricate(:group) }
   fab!(:topic) { Fabricate(:topic) }
@@ -18,13 +18,6 @@ RSpec.describe "AI Composer helper", type: :system, js: true do
   let(:topic_page) { PageObjects::Pages::Topic.new }
   let(:post_ai_helper) { PageObjects::Components::AIHelperPostOptions.new }
 
-  let(:explain_response) { <<~STRING }
-    In this context, \"pie\" refers to a baked dessert typically consisting of a pastry crust and filling.
-    The person states they enjoy eating pie, considering it a good dessert. They note that some people wastefully
-    throw pie at others, but the person themselves chooses to eat the pie rather than throwing it. Overall, \"pie\"
-    is being used to refer the the baked dessert food item.
-  STRING
-
   before do
     Group.find_by(id: Group::AUTO_GROUPS[:admins]).add(user)
     SiteSetting.composer_ai_helper_enabled = true
@@ -56,18 +49,23 @@ RSpec.describe "AI Composer helper", type: :system, js: true do
     end
 
     context "when using explain mode" do
-      skip "TODO: Fix explain mode option not appearing in spec" do
       let(:mode) { CompletionPrompt::EXPLAIN }
 
+      let(:explain_response) { <<~STRING }
+        In this context, \"pie\" refers to a baked dessert typically consisting of a pastry crust and filling.
+        The person states they enjoy eating pie, considering it a good dessert. They note that some people wastefully
+        throw pie at others, but the person themselves chooses to eat the pie rather than throwing it. Overall, \"pie\"
+        is being used to refer the the baked dessert food item.
+      STRING
+
+      skip "TODO: Fix explain option stuck in loading in test" do
         it "shows an explanation of the selected text" do
           select_post_text(post)
           post_ai_helper.click_ai_button
 
           DiscourseAi::Completions::Llm.with_prepared_responses([explain_response]) do
             post_ai_helper.select_helper_model(mode)
-
             wait_for { post_ai_helper.suggestion_value == explain_response }
-
             expect(post_ai_helper.suggestion_value).to eq(explain_response)
           end
         end
@@ -75,7 +73,6 @@ RSpec.describe "AI Composer helper", type: :system, js: true do
     end
 
     context "when using translate mode" do
-      skip "TODO: Fix WebMock request for translate mode not working" do
       let(:mode) { CompletionPrompt::TRANSLATE }
 
       let(:translated_input) { "The rain in Spain, stays mainly in the Plane." }
@@ -94,7 +91,6 @@ RSpec.describe "AI Composer helper", type: :system, js: true do
       end
     end
   end
-  end
 
   context "when AI helper is disabled" do
     before { SiteSetting.composer_ai_helper_enabled = false }

spec/system/page_objects/components/ai_helper_post_options.rb

@@ -23,7 +23,7 @@ module PageObjects
       end
 
       def suggestion_value
-        find(SUGGESTION_SELECTOR).text
+        find("#{SUGGESTION_SELECTOR}__text").text
       end
 
       def has_post_ai_helper?