WIP: Add streaming to composer helper

Author: keegangeorge

app/controllers/discourse_ai/ai_helper/assistant_controller.rb

@@ -110,26 +110,43 @@ def suggest_thumbnails(input)
       end
 
       def stream_suggestion
-        post_id = get_post_param!
         text = get_text_param!
-        post = Post.includes(:topic).find_by(id: post_id)
+
+        location = params[:location]
+        raise Discourse::InvalidParameters.new(:location) if !location
+
         prompt = CompletionPrompt.find_by(id: params[:mode])
 
         raise Discourse::InvalidParameters.new(:mode) if !prompt || !prompt.enabled?
-        raise Discourse::InvalidParameters.new(:post_id) unless post
+        return suggest_thumbnails(input) if prompt.id == CompletionPrompt::ILLUSTRATE_POST
 
         if prompt.id == CompletionPrompt::CUSTOM_PROMPT
           raise Discourse::InvalidParameters.new(:custom_prompt) if params[:custom_prompt].blank?
         end
 
-        Jobs.enqueue(
-          :stream_post_helper,
-          post_id: post.id,
-          user_id: current_user.id,
-          text: text,
-          prompt: prompt.name,
-          custom_prompt: params[:custom_prompt],
-        )
+        if location == "composer"
+          Jobs.enqueue(
+            :stream_composer_helper,
+            user_id: current_user.id,
+            text: text,
+            prompt: prompt.name,
+            custom_prompt: params[:custom_prompt],
+          )
+        else
+          post_id = get_post_param!
+          post = Post.includes(:topic).find_by(id: post_id)
+
+          raise Discourse::InvalidParameters.new(:post_id) unless post
+
+          Jobs.enqueue(
+            :stream_post_helper,
+            post_id: post.id,
+            user_id: current_user.id,
+            text: text,
+            prompt: prompt.name,
+            custom_prompt: params[:custom_prompt],
+          )
+        end
 
         render json: { success: true }, status: 200
       rescue DiscourseAi::Completions::Endpoints::Base::CompletionFailed

app/jobs/regular/stream_composer_helper.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Jobs
+  class StreamComposerHelper < ::Jobs::Base
+    sidekiq_options retry: false
+    # TODO handle force_default_locale stuff
+    def execute(args)
+      return unless args[:prompt]
+      return unless user = User.find_by(id: args[:user_id])
+      return unless args[:text]
+
+      prompt = CompletionPrompt.enabled_by_name(args[:prompt])
+
+      if prompt.id == CompletionPrompt::CUSTOM_PROMPT
+        prompt.custom_instruction = args[:custom_prompt]
+      end
+
+      DiscourseAi::AiHelper::Assistant.new.stream_prompt(
+        prompt,
+        args[:text],
+        user,
+        "/discourse-ai/ai-helper/stream_suggestion",
+      )
+    end
+  end
+end

assets/javascripts/discourse/components/modal/diff-modal.gjs

@@ -1,18 +1,22 @@
 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 { service } from "@ember/service";
 import { htmlSafe } from "@ember/template";
 import CookText from "discourse/components/cook-text";
 import DButton from "discourse/components/d-button";
 import DModal from "discourse/components/d-modal";
 import { ajax } from "discourse/lib/ajax";
 import { popupAjaxError } from "discourse/lib/ajax-error";
+import { bind } from "discourse/lib/decorators";
 import { i18n } from "discourse-i18n";
 import AiIndicatorWave from "../ai-indicator-wave";
 
 export default class ModalDiffModal extends Component {
   @service currentUser;
+  @service messageBus;
 
   @tracked loading = false;
   @tracked diff;
@@ -23,23 +27,47 @@ export default class ModalDiffModal extends Component {
     this.suggestChanges();
   }
 
+  @bind
+  subscribe() {
+    console.log("subscribe called");
+    const channel = "/discourse-ai/ai-helper/stream_suggestion";
+    this.messageBus.subscribe(channel, this.updateResult);
+  }
+
+  @bind
+  unsubscribe() {
+    console.log("unsubscribe called");
+    const channel = "/discourse-ai/ai-helper/stream_suggestion";
+    this.messageBus.subscribe(channel, this.updateResult);
+  }
+
+  @action
+  async updateResult(result) {
+    console.log("updateResult called");
+    console.log("result", result);
+
+    this.diff = result.diff;
+    this.suggestion = result.result;
+  }
+
   @action
   async suggestChanges() {
     this.loading = true;
 
     try {
-      const suggestion = await ajax("/discourse-ai/ai-helper/suggest", {
+      return await ajax("/discourse-ai/ai-helper/stream_suggestion", {
         method: "POST",
         data: {
+          location: "composer",
           mode: this.args.model.mode,
           text: this.args.model.selectedText,
           custom_prompt: this.args.model.customPromptValue,
           force_default_locale: true,
         },
       });
 
-      this.diff = suggestion.diff;
-      this.suggestion = suggestion.suggestions[0];
+      // this.diff = suggestion.diff;
+      // this.suggestion = suggestion.suggestions[0];
     } catch (e) {
       popupAjaxError(e);
     } finally {
@@ -71,17 +99,23 @@ export default class ModalDiffModal extends Component {
             <CookText @rawText={{@model.selectedText}} />
           </div>
         {{else}}
-          {{#if this.diff}}
-            {{htmlSafe this.diff}}
-          {{else}}
-            <div class="composer-ai-helper-modal__old-value">
-              {{@model.selectedText}}
-            </div>
-
-            <div class="composer-ai-helper-modal__new-value">
-              {{this.suggestion}}
-            </div>
-          {{/if}}
+          <div
+            class="composer-ai-helper-modal__suggestion"
+            {{didInsert this.subscribe}}
+            {{willDestroy this.unsubscribe}}
+          >
+            {{#if this.diff}}
+              {{htmlSafe this.diff}}
+            {{else}}
+              <div class="composer-ai-helper-modal__old-value">
+                {{@model.selectedText}}
+              </div>
+
+              <div class="composer-ai-helper-modal__new-value">
+                {{this.suggestion}}
+              </div>
+            {{/if}}
+          </div>
         {{/if}}
 
       </:body>

lib/ai_helper/assistant.rb

@@ -161,23 +161,28 @@ def generate_and_send_prompt(completion_prompt, input, user, force_default_local
       end
 
       def stream_prompt(completion_prompt, input, user, channel)
+        streamed_diff = +""
         streamed_result = +""
         start = Time.now
 
         generate_prompt(completion_prompt, input, user) do |partial_response, cancel_function|
           streamed_result << partial_response
 
+          streamed_diff = parse_diff(input, partial_response) if completion_prompt.diff?
+
           # Throttle the updates
           if (Time.now - start > 0.5) || Rails.env.test?
-            payload = { result: sanitize_result(streamed_result), done: false }
+            payload = { result: sanitize_result(streamed_result), diff: streamed_diff, done: false }
             publish_update(channel, payload, user)
             start = Time.now
           end
         end
 
+        final_diff = parse_diff(input, streamed_result) if completion_prompt.diff?
+
         sanitized_result = sanitize_result(streamed_result)
         if sanitized_result.present?
-          publish_update(channel, { result: sanitized_result, done: true }, user)
+          publish_update(channel, { result: sanitized_result, diff: final_diff, done: true }, user)
         end
       end