FEATURE: forum researcher persona for deep research

assets/javascripts/discourse/lib/ai-streamer/progress-handlers.js

@@ -2,7 +2,7 @@ import { later } from "@ember/runloop";
 import PostUpdater from "./updaters/post-updater";
 
 const PROGRESS_INTERVAL = 40;
-const GIVE_UP_INTERVAL = 60000;
+const GIVE_UP_INTERVAL = 600000; // 10 minutes which is our max thinking time for now
 export const MIN_LETTERS_PER_INTERVAL = 6;
 const MAX_FLUSH_TIME = 800;
 

config/locales/server.en.yml

@@ -296,6 +296,9 @@ en:
         designer:
           name: Designer
           description: "AI Bot specialized in generating and editing images"
+        forum_researcher:
+          name: Forum Researcher
+          description: "AI Bot specialized in deep research for the forum"
         sql_helper:
           name: SQL Helper
           description: "AI Bot specialized in helping craft SQL queries on this Discourse instance"
@@ -327,6 +330,16 @@ en:
       summarizing: "Summarizing topic"
       searching: "Searching for: '%{query}'"
       tool_options:
+        researcher:
+          max_results:
+            name: "Maximum number of results"
+            description: "Maximum number of results to include in a filter"
+          include_private:
+            name: "Include private"
+            description: "Include private topics in the filters"
+          max_tokens_per_post:
+            name: "Maximum tokens per post"
+            description: "Maximum number of tokens to use for each post in the filter"
         create_artifact:
           creator_llm:
             name: "LLM"
@@ -385,6 +398,7 @@ en:
         javascript_evaluator: "Evaluate JavaScript"
         create_image: "Creating image"
         edit_image: "Editing image"
+        researcher: "Researching"
       tool_help:
         read_artifact: "Read a web artifact using the AI Bot"
         update_artifact: "Update a web artifact using the AI Bot"
@@ -411,6 +425,7 @@ en:
         dall_e: "Generate image using DALL-E 3"
         search_meta_discourse: "Search Meta Discourse"
         javascript_evaluator: "Evaluate JavaScript"
+        researcher: "Research forum information using the AI Bot"
       tool_description:
         read_artifact: "Read a web artifact using the AI Bot"
         update_artifact: "Updated a web artifact using the AI Bot"
@@ -445,6 +460,12 @@ en:
           other: "Found %{count} <a href='%{url}'>results</a> for '%{query}'"
         setting_context: "Reading context for: %{setting_name}"
         schema: "%{tables}"
+        researcher_dry_run:
+          one: "Proposed research: %{goals}\n\nFound %{count} result for '%{filter}'"
+          other: "Proposed research: %{goals}\n\nFound %{count} result for '%{filter}'"
+        researcher:
+          one: "Researching: %{goals}\n\nFound %{count} result for '%{filter}'"
+          other: "Researching: %{goals}\n\nFound %{count} result for '%{filter}'"
         search_settings:
           one: "Found %{count} result for '%{query}'"
           other: "Found %{count} results for '%{query}'"

db/fixtures/personas/603_ai_personas.rb

@@ -33,7 +33,7 @@ def from_setting(setting_name)
       persona.allowed_group_ids = [Group::AUTO_GROUPS[:trust_level_0]]
     end
 
-    persona.enabled = !summarization_personas.include?(persona_class)
+    persona.enabled = persona_class.default_enabled
     persona.priority = true if persona_class == DiscourseAi::Personas::General
   end
 

lib/ai_bot/chat_streamer.rb

@@ -6,16 +6,23 @@
 module DiscourseAi
   module AiBot
     class ChatStreamer
-      attr_accessor :cancel
       attr_reader :reply,
                   :guardian,
                   :thread_id,
                   :force_thread,
                   :in_reply_to_id,
                   :channel,
-                  :cancelled
-
-      def initialize(message:, channel:, guardian:, thread_id:, in_reply_to_id:, force_thread:)
+                  :cancel_manager
+
+      def initialize(
+        message:,
+        channel:,
+        guardian:,
+        thread_id:,
+        in_reply_to_id:,
+        force_thread:,
+        cancel_manager: nil
+      )
         @message = message
         @channel = channel
         @guardian = guardian
@@ -35,6 +42,8 @@ def initialize(message:, channel:, guardian:, thread_id:, in_reply_to_id:, force
             guardian: guardian,
             thread_id: thread_id,
           )
+
+        @cancel_manager = cancel_manager
       end
 
       def <<(partial)
@@ -111,8 +120,7 @@ def run
 
           streaming = ChatSDK::Message.stream(message_id: reply.id, raw: buffer, guardian: guardian)
           if !streaming
-            cancel.call
-            @cancelled = true
+            @cancel_manager.cancel! if @cancel_manager
           end
         end
       end

lib/ai_bot/playground.rb

@@ -331,6 +331,7 @@ def reply_to_chat_message(message, channel, context_post_ids)
               ),
             user: message.user,
             skip_tool_details: true,
+            cancel_manager: DiscourseAi::Completions::CancelManager.new,
           )
 
         reply = nil
@@ -347,15 +348,14 @@ def reply_to_chat_message(message, channel, context_post_ids)
             thread_id: message.thread_id,
             in_reply_to_id: in_reply_to_id,
             force_thread: force_thread,
+            cancel_manager: context.cancel_manager,
           )
 
         new_prompts =
-          bot.reply(context) do |partial, cancel, placeholder, type|
+          bot.reply(context) do |partial, placeholder, type|
             # no support for tools or thinking by design
             next if type == :thinking || type == :tool_details || type == :partial_tool
-            streamer.cancel = cancel
             streamer << partial
-            break if streamer.cancelled
           end
 
         reply = streamer.reply
@@ -383,6 +383,7 @@ def reply_to(
         auto_set_title: true,
         silent_mode: false,
         feature_name: nil,
+        cancel_manager: nil,
         &blk
       )
         # this is a multithreading issue
@@ -471,16 +472,26 @@ def reply_to(
 
           redis_stream_key = "gpt_cancel:#{reply_post.id}"
           Discourse.redis.setex(redis_stream_key, MAX_STREAM_DELAY_SECONDS, 1)
+
+          cancel_manager ||= DiscourseAi::Completions::CancelManager.new
+          context.cancel_manager = cancel_manager
+          context
+            .cancel_manager
+            .start_monitor(delay: 0.2) do
+              context.cancel_manager.cancel! if !Discourse.redis.get(redis_stream_key)
+            end
+
+          context.cancel_manager.add_callback(
+            lambda { reply_post.update!(raw: reply, cooked: PrettyText.cook(reply)) },
+          )
         end
 
         context.skip_tool_details ||= !bot.persona.class.tool_details
-
         post_streamer = PostStreamer.new(delay: Rails.env.test? ? 0 : 0.5) if stream_reply
-
         started_thinking = false
 
         new_custom_prompts =
-          bot.reply(context) do |partial, cancel, placeholder, type|
+          bot.reply(context) do |partial, placeholder, type|
             if type == :thinking && !started_thinking
               reply << "<details><summary>#{I18n.t("discourse_ai.ai_bot.thinking")}</summary>"
               started_thinking = true
@@ -499,15 +510,6 @@ def reply_to(
               blk.call(partial)
             end
 
-            if stream_reply && !Discourse.redis.get(redis_stream_key)
-              cancel&.call
-              reply_post.update!(raw: reply, cooked: PrettyText.cook(reply))
-              # we do not break out, cause if we do
-              # we will not get results from bot
-              # leading to broken context
-              # we need to trust it to cancel at the endpoint
-            end
-
             if post_streamer
               post_streamer.run_later do
                 Discourse.redis.expire(redis_stream_key, MAX_STREAM_DELAY_SECONDS)
@@ -568,6 +570,8 @@ def reply_to(
         end
         raise e
       ensure
+        context.cancel_manager.stop_monitor if context&.cancel_manager
+
         # since we are skipping validations and jobs we
         # may need to fix participant count
         if reply_post && reply_post.topic && reply_post.topic.private_message? &&
@@ -649,7 +653,7 @@ def publish_update(bot_reply_post, payload)
           payload,
           user_ids: bot_reply_post.topic.allowed_user_ids,
           max_backlog_size: 2,
-          max_backlog_age: 60,
+          max_backlog_age: MAX_STREAM_DELAY_SECONDS,
         )
       end
     end

lib/completions/cancel_manager.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+# special object that can be used to cancel completions and http requests
+module DiscourseAi
+  module Completions
+    class CancelManager
+      attr_reader :cancelled
+      attr_reader :callbacks
+
+      def initialize
+        @cancelled = false
+        @callbacks = Concurrent::Array.new
+        @mutex = Mutex.new
+        @monitor_thread = nil
+      end
+
+      def monitor_thread
+        @mutex.synchronize { @monitor_thread }
+      end
+
+      def start_monitor(delay: 0.5, &block)
+        @mutex.synchronize do
+          raise "Already monitoring" if @monitor_thread
+          raise "Expected a block" if !block
+
+          db = RailsMultisite::ConnectionManagement.current_db
+          @stop_monitor = false
+
+          @monitor_thread =
+            Thread.new do
+              begin
+                loop do
+                  done = false
+                  @mutex.synchronize { done = true if @stop_monitor }
+                  break if done
+                  sleep delay
+                  @mutex.synchronize { done = true if @stop_monitor }
+                  @mutex.synchronize { done = true if cancelled? }
+                  break if done
+
+                  should_cancel = false
+                  RailsMultisite::ConnectionManagement.with_connection(db) do
+                    should_cancel = block.call
+                  end
+
+                  @mutex.synchronize { cancel! if should_cancel }
+
+                  break if cancelled?
+                end
+              ensure
+                @mutex.synchronize { @monitor_thread = nil }
+              end
+            end
+        end
+      end
+
+      def stop_monitor
+        monitor_thread = nil
+
+        @mutex.synchronize { monitor_thread = @monitor_thread }
+
+        if monitor_thread
+          @mutex.synchronize { @stop_monitor = true }
+          # so we do not deadlock
+          monitor_thread.wakeup
+          monitor_thread.join(2)
+          # should not happen
+          if monitor_thread.alive?
+            Rails.logger.warn("DiscourseAI: CancelManager monitor thread did not stop in time")
+            monitor_thread.kill if monitor_thread.alive?
+          end
+          @monitor_thread = nil
+        end
+      end
+
+      def cancelled?
+        @cancelled
+      end
+
+      def add_callback(cb)
+        @callbacks << cb
+      end
+
+      def remove_callback(cb)
+        @callbacks.delete(cb)
+      end
+
+      def cancel!
+        @cancelled = true
+        monitor_thread = @monitor_thread
+        if monitor_thread && monitor_thread != Thread.current
+          monitor_thread.wakeup
+          monitor_thread.join(2)
+          if monitor_thread.alive?
+            Rails.logger.warn("DiscourseAI: CancelManager monitor thread did not stop in time")
+            monitor_thread.kill if monitor_thread.alive?
+          end
+        end
+        @callbacks.each do |cb|
+          begin
+            cb.call
+          rescue StandardError
+            # ignore cause this may have already been cancelled
+          end
+        end
+      end
+    end
+  end
+end