From 6aaf1f002e6015589381be051638fe5e5808a098 Mon Sep 17 00:00:00 2001 From: Keegan George Date: Tue, 12 Dec 2023 09:28:39 -0800 Subject: [PATCH] FEATURE: Add streaming to post AI helper's explain option (#344) Co-authored-by: Rafael dos Santos Silva Co-authored-by: Roman Rizzi --- .../ai_helper/assistant_controller.rb | 14 ++- app/jobs/regular/stream_post_helper.rb | 35 ++++++ .../components/ai-helper-loading.gjs | 19 +++ .../after-d-editor/ai-helper-context-menu.hbs | 13 +-- .../ai-helper-options-menu.gjs | 108 ++++++++++++------ .../modules/ai-helper/common/ai-helper.scss | 13 ++- config/locales/client.en.yml | 1 + lib/ai_helper/assistant.rb | 57 ++++++++- lib/ai_helper/topic_helper.rb | 35 ------ spec/jobs/regular/stream_post_helper.rb | 83 ++++++++++++++ spec/system/ai_helper/ai_post_helper_spec.rb | 42 +++---- .../components/ai_helper_post_options.rb | 2 +- 12 files changed, 302 insertions(+), 120 deletions(-) create mode 100644 app/jobs/regular/stream_post_helper.rb create mode 100644 assets/javascripts/discourse/components/ai-helper-loading.gjs delete mode 100644 lib/ai_helper/topic_helper.rb create mode 100644 spec/jobs/regular/stream_post_helper.rb diff --git a/app/controllers/discourse_ai/ai_helper/assistant_controller.rb b/app/controllers/discourse_ai/ai_helper/assistant_controller.rb index eea254ec..ec9863fe 100644 --- a/app/controllers/discourse_ai/ai_helper/assistant_controller.rb +++ b/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: - DiscourseAi::AiHelper::TopicHelper.new(current_user).explain( - term_to_explain, - post, - ), - status: 200 + Jobs.enqueue( + :stream_post_helper, + post_id: post.id, + user_id: current_user.id, + term_to_explain: term_to_explain, + ) + + 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 diff --git a/app/jobs/regular/stream_post_helper.rb b/app/jobs/regular/stream_post_helper.rb new file mode 100644 index 00000000..9b103d06 --- /dev/null +++ b/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 + #{args[:term_to_explain]} + #{post.raw} + #{topic.title} + #{reply_to ? "#{reply_to.raw}" : nil} + TEXT + + DiscourseAi::AiHelper::Assistant.new.stream_prompt( + prompt, + input, + user, + "/discourse-ai/ai-helper/explain/#{post.id}", + ) + end + end +end diff --git a/assets/javascripts/discourse/components/ai-helper-loading.gjs b/assets/javascripts/discourse/components/ai-helper-loading.gjs new file mode 100644 index 00000000..730be41d --- /dev/null +++ b/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 = ; + +export default AiHelperLoading; 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 f61a6828..b63cad7a 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 @@ -53,18 +53,7 @@ {{else if (eq this.menuState this.CONTEXT_MENU_STATES.loading)}} -
-
- - {{i18n "discourse_ai.ai_helper.context_menu.loading"}} - - -
+ {{else if (eq this.menuState this.CONTEXT_MENU_STATES.review)}}
    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 753b4513..81e779af 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 @@ -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,15 +100,17 @@ export default class AIHelperOptionsMenu extends Component { }); } - this._activeAIRequest - .then(({ suggestions }) => { - this.suggestion = suggestions[0]; - }) - .catch(popupAjaxError) - .finally(() => { - this.loading = false; - this.menuState = this.MENU_STATES.result; - }); + if (option.name !== "Explain") { + this._activeAIRequest + .then(({ suggestions }) => { + this.suggestion = suggestions[0]; + }) + .catch(popupAjaxError) + .finally(() => { + this.loading = false; + this.menuState = this.MENU_STATES.result; + }); + } return this._activeAIRequest; } @@ -154,30 +185,35 @@ export default class AIHelperOptionsMenu extends Component { {{else if (eq this.menuState this.MENU_STATES.loading)}} -
    -
    - - {{i18n "discourse_ai.ai_helper.context_menu.loading"}} - - -
    + {{else if (eq this.menuState this.MENU_STATES.result)}} -
    -
    - {{this.suggestion}} -
    - +
    + {{#if this.suggestion}} +
    + {{this.suggestion}} +
    + + + + + {{else}} + + {{/if}}
    {{/if}}
    diff --git a/assets/stylesheets/modules/ai-helper/common/ai-helper.scss b/assets/stylesheets/modules/ai-helper/common/ai-helper.scss index bf543bb0..530c8cf0 100644 --- a/assets/stylesheets/modules/ai-helper/common/ai-helper.scss +++ b/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%; + } + } } } diff --git a/config/locales/client.en.yml b/config/locales/client.en.yml index 129feb4f..29ef2bd2 100644 --- a/config/locales/client.en.yml +++ b/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:" diff --git a/lib/ai_helper/assistant.rb b/lib/ai_helper/assistant.rb index 8444e225..139ceacb 100644 --- a/lib/ai_helper/assistant.rb +++ b/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[ + + + + + + + + + + + + + ] + + 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" diff --git a/lib/ai_helper/topic_helper.rb b/lib/ai_helper/topic_helper.rb deleted file mode 100644 index 2ccad788..00000000 --- a/lib/ai_helper/topic_helper.rb +++ /dev/null @@ -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_to_explain} - #{post.raw} - #{topic.title} - #{reply_to ? "#{reply_to.raw}" : nil} - TEXT - - DiscourseAi::AiHelper::Assistant.new.generate_and_send_prompt(prompt, input, user) - end - - private - - attr_reader :user - end - end -end diff --git a/spec/jobs/regular/stream_post_helper.rb b/spec/jobs/regular/stream_post_helper.rb new file mode 100644 index 00000000..dab2dd0e --- /dev/null +++ b/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 diff --git a/spec/system/ai_helper/ai_post_helper_spec.rb b/spec/system/ai_helper/ai_post_helper_spec.rb index 4569554a..fab45e67 100644 --- a/spec/system/ai_helper/ai_post_helper_spec.rb +++ b/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(: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,22 +73,20 @@ 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(:mode) { CompletionPrompt::TRANSLATE } - let(:translated_input) { "The rain in Spain, stays mainly in the Plane." } + let(:translated_input) { "The rain in Spain, stays mainly in the Plane." } - it "shows a translation of the selected text" do - select_post_text(post_2) - post_ai_helper.click_ai_button + it "shows a translation of the selected text" do + select_post_text(post_2) + post_ai_helper.click_ai_button - DiscourseAi::Completions::Llm.with_prepared_responses([translated_input]) do - post_ai_helper.select_helper_model(mode) + DiscourseAi::Completions::Llm.with_prepared_responses([translated_input]) do + post_ai_helper.select_helper_model(mode) - wait_for { post_ai_helper.suggestion_value == translated_input } + wait_for { post_ai_helper.suggestion_value == translated_input } - expect(post_ai_helper.suggestion_value).to eq(translated_input) - end + expect(post_ai_helper.suggestion_value).to eq(translated_input) end end end diff --git a/spec/system/page_objects/components/ai_helper_post_options.rb b/spec/system/page_objects/components/ai_helper_post_options.rb index 73f8fe5b..942d5b04 100644 --- a/spec/system/page_objects/components/ai_helper_post_options.rb +++ b/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?