FEATURE: display more places where AI is used / Chat streamer (#1278)

* FEATURE: display more places where AI is used

- Usage was not showing automation or image caption in llm list.
- Also: FIX - reasoning models would time out incorrectly after 60 seconds (raised to 10 minutes)

* correct enum not to enumerate non configured models

* FEATURE: implement chat streamer

This implements a basic chat streamer, it provides 2 things:

1. Gives feedback to the user when LLM is generating
2. Streams stuff much more efficiently to client (given it may take 100ms or so per call to update chat)

Author: SamSaffron

assets/javascripts/discourse/components/ai-llms-list-editor.gjs

@@ -116,6 +116,10 @@ export default class AiLlmsListEditor extends Component {
       return i18n("discourse_ai.llms.usage.ai_persona", {
         persona: usage.name,
       });
+    } else if (usage.type === "automation") {
+      return i18n("discourse_ai.llms.usage.automation", {
+        name: usage.name,
+      });
     } else {
       return i18n("discourse_ai.llms.usage." + usage.type);
     }

assets/stylesheets/modules/llms/common/ai-llms-editor.scss

@@ -117,6 +117,7 @@
   list-style: none;
   margin: 0.5em 0 0 0;
   display: flex;
+  flex-wrap: wrap;
 
   li {
     font-size: var(--font-down-2);
@@ -125,6 +126,7 @@
     border: 1px solid var(--primary-low);
     padding: 1px 3px;
     margin-right: 0.5em;
+    margin-bottom: 0.5em;
   }
 }
 

config/locales/client.en.yml

@@ -439,10 +439,12 @@ en:
         usage:
           ai_bot: "AI bot"
           ai_helper: "Helper"
+          ai_helper_image_caption: "Image caption"
           ai_persona: "Persona (%{persona})"
           ai_summarization: "Summarize"
           ai_embeddings_semantic_search: "AI search"
           ai_spam: "Spam"
+          automation: "Automation (%{name})"
         in_use_warning:
           one: "This model is currently used by %{settings}. If misconfigured, the feature won't work as expected."
           other: "This model is currently used by the following: %{settings}. If misconfigured, features won't work as expected. "

config/settings.yml

@@ -312,11 +312,6 @@ discourse_ai:
     default: "1|2" # 1: admins, 2: moderators
     allow_any: false
     refresh: true
-  ai_bot_enabled_chat_bots: # TODO(roman): Deprecated. Remove by Sept 2024
-    type: list
-    default: "gpt-3.5-turbo"
-    hidden: true
-    choices: "DiscourseAi::Configuration::LlmEnumerator.available_ai_bots"
   ai_bot_add_to_header:
     default: true
     client: true

db/migrate/20250424035234_remove_old_settings.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+class RemoveOldSettings < ActiveRecord::Migration[7.2]
+  def up
+    execute <<~SQL
+      DELETE FROM site_settings
+      WHERE name IN ('ai_bot_enabled_chat_bots')
+    SQL
+  end
+
+  def down
+    raise ActiveRecord::IrreversibleMigration
+  end
+end

lib/ai_bot/chat_streamer.rb

@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+#
+# Chat streaming APIs are a bit slow, this ensures we properly buffer results
+# and stream as quickly as possible.
+
+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:)
+        @message = message
+        @channel = channel
+        @guardian = guardian
+        @thread_id = thread_id
+        @force_thread = force_thread
+        @in_reply_to_id = in_reply_to_id
+
+        @queue = Queue.new
+
+        db = RailsMultisite::ConnectionManagement.current_db
+        @worker_thread =
+          Thread.new { RailsMultisite::ConnectionManagement.with_connection(db) { run } }
+
+        @client_id =
+          ChatSDK::Channel.start_reply(
+            channel_id: message.chat_channel_id,
+            guardian: guardian,
+            thread_id: thread_id,
+          )
+      end
+
+      def <<(partial)
+        return if partial.to_s.empty?
+
+        if @client_id
+          ChatSDK::Channel.stop_reply(
+            channel_id: @message.chat_channel_id,
+            client_id: @client_id,
+            guardian: @guardian,
+            thread_id: @thread_id,
+          )
+          @client_id = nil
+        end
+
+        if @reply
+          @queue << partial
+        else
+          create_reply(partial)
+        end
+      end
+
+      def create_reply(message)
+        @reply =
+          ChatSDK::Message.create(
+            raw: message,
+            channel_id: channel.id,
+            guardian: guardian,
+            force_thread: force_thread,
+            in_reply_to_id: in_reply_to_id,
+            enforce_membership: !channel.direct_message_channel?,
+          )
+
+        ChatSDK::Message.start_stream(message_id: @reply.id, guardian: @guardian)
+
+        if trailing = message.scan(/\s*\z/).first
+          @queue << trailing
+        end
+      end
+
+      def done
+        @queue << :done
+        @worker_thread.join
+        ChatSDK::Message.stop_stream(message_id: @reply.id, guardian: @guardian)
+        @reply
+      end
+
+      private
+
+      def run
+        done = false
+        while !done
+          buffer = +""
+          popped = @queue.pop
+          break if popped == :done
+
+          buffer << popped
+
+          begin
+            while true
+              popped = @queue.pop(true)
+              if popped == :done
+                done = true
+                break
+              end
+              buffer << popped
+            end
+          rescue ThreadError
+          end
+
+          streaming = ChatSDK::Message.stream(message_id: reply.id, raw: buffer, guardian: guardian)
+          if !streaming
+            cancel.call
+            @cancelled = true
+          end
+        end
+      end
+    end
+  end
+end

lib/ai_bot/playground.rb

@@ -5,6 +5,9 @@ module AiBot
     class Playground
       BYPASS_AI_REPLY_CUSTOM_FIELD = "discourse_ai_bypass_ai_reply"
       BOT_USER_PREF_ID_CUSTOM_FIELD = "discourse_ai_bot_user_pref_id"
+      # 10 minutes is enough for vast majority of cases
+      # there is a small chance that some reasoning models may take longer
+      MAX_STREAM_DELAY_SECONDS = 600
 
       attr_reader :bot
 
@@ -334,42 +337,38 @@ def reply_to_chat_message(message, channel, context_post_ids)
         force_thread = message.thread_id.nil? && channel.direct_message_channel?
         in_reply_to_id = channel.direct_message_channel? ? message.id : nil
 
+        streamer =
+          ChatStreamer.new(
+            message: message,
+            channel: channel,
+            guardian: guardian,
+            thread_id: message.thread_id,
+            in_reply_to_id: in_reply_to_id,
+            force_thread: force_thread,
+          )
+
         new_prompts =
           bot.reply(context) do |partial, cancel, placeholder, type|
             # no support for tools or thinking by design
             next if type == :thinking || type == :tool_details || type == :partial_tool
-            if !reply
-              # just eat all leading spaces we can not create the message
-              next if partial.blank?
-              reply =
-                ChatSDK::Message.create(
-                  raw: partial,
-                  thread_id: message.thread_id,
-                  channel_id: channel.id,
-                  guardian: guardian,
-                  in_reply_to_id: in_reply_to_id,
-                  force_thread: force_thread,
-                  enforce_membership: !channel.direct_message_channel?,
-                )
-              ChatSDK::Message.start_stream(message_id: reply.id, guardian: guardian)
-            else
-              streaming =
-                ChatSDK::Message.stream(message_id: reply.id, raw: partial, guardian: guardian)
-
-              if !streaming
-                cancel&.call
-                break
-              end
-            end
+            streamer.cancel = cancel
+            streamer << partial
+            break if streamer.cancelled
           end
 
-        if new_prompts.length > 1 && reply.id
+        reply = streamer.reply
+        if new_prompts.length > 1 && reply
           ChatMessageCustomPrompt.create!(message_id: reply.id, custom_prompt: new_prompts)
         end
 
-        ChatSDK::Message.stop_stream(message_id: reply.id, guardian: guardian) if reply
+        if streamer
+          streamer.done
+          streamer = nil
+        end
 
         reply
+      ensure
+        streamer.done if streamer
       end
 
       def reply_to(
@@ -464,7 +463,7 @@ def reply_to(
           publish_update(reply_post, { raw: reply_post.cooked })
 
           redis_stream_key = "gpt_cancel:#{reply_post.id}"
-          Discourse.redis.setex(redis_stream_key, 60, 1)
+          Discourse.redis.setex(redis_stream_key, MAX_STREAM_DELAY_SECONDS, 1)
         end
 
         context.skip_tool_details ||= !bot.persona.class.tool_details
@@ -504,7 +503,7 @@ def reply_to(
 
             if post_streamer
               post_streamer.run_later do
-                Discourse.redis.expire(redis_stream_key, 60)
+                Discourse.redis.expire(redis_stream_key, MAX_STREAM_DELAY_SECONDS)
                 publish_update(reply_post, { raw: raw })
               end
             end

lib/configuration/llm_enumerator.rb

@@ -13,16 +13,22 @@ def self.global_usage
             .where("enabled_chat_bot = ?", true)
             .pluck(:id)
             .each { |llm_id| rval[llm_id] << { type: :ai_bot } }
-
-          AiPersona
-            .where("force_default_llm = ?", true)
-            .pluck(:default_llm_id, :name, :id)
-            .each { |llm_id, name, id| rval[llm_id] << { type: :ai_persona, name: name, id: id } }
         end
 
+        # this is unconditional, so it is clear that we always signal configuration
+        AiPersona
+          .where("default_llm_id IS NOT NULL")
+          .pluck(:default_llm_id, :name, :id)
+          .each { |llm_id, name, id| rval[llm_id] << { type: :ai_persona, name: name, id: id } }
+
         if SiteSetting.ai_helper_enabled
           model_id = SiteSetting.ai_helper_model.split(":").last.to_i
-          rval[model_id] << { type: :ai_helper }
+          rval[model_id] << { type: :ai_helper } if model_id != 0
+        end
+
+        if SiteSetting.ai_helper_image_caption_model
+          model_id = SiteSetting.ai_helper_image_caption_model.split(":").last.to_i
+          rval[model_id] << { type: :ai_helper_image_caption } if model_id != 0
         end
 
         if SiteSetting.ai_summarization_enabled
@@ -42,6 +48,25 @@ def self.global_usage
           rval[model_id] << { type: :ai_spam }
         end
 
+        if defined?(DiscourseAutomation::Automation)
+          DiscourseAutomation::Automation
+            .joins(:fields)
+            .where(script: %w[llm_report llm_triage])
+            .where("discourse_automation_fields.name = ?", "model")
+            .pluck(
+              "metadata ->> 'value', discourse_automation_automations.name, discourse_automation_automations.id",
+            )
+            .each do |model_text, name, id|
+              next if model_text.blank?
+              model_id = model_text.split("custom:").last.to_i
+              if model_id.present?
+                if model_text =~ /custom:(\d+)/
+                  rval[model_id] << { type: :automation, name: name, id: id }
+                end
+              end
+            end
+        end
+
         rval
       end
 
@@ -85,45 +110,6 @@ def self.values(allowed_seeded_llms: nil)
         values.each { |value_h| value_h[:value] = "custom:#{value_h[:value]}" }
         values
       end
-
-      # TODO(roman): Deprecated. Remove by Sept 2024
-      def self.old_summarization_options
-        %w[
-          gpt-4
-          gpt-4-32k
-          gpt-4-turbo
-          gpt-4o
-          gpt-3.5-turbo
-          gpt-3.5-turbo-16k
-          gemini-pro
-          gemini-1.5-pro
-          gemini-1.5-flash
-          claude-2
-          claude-instant-1
-          claude-3-haiku
-          claude-3-sonnet
-          claude-3-opus
-          mistralai/Mixtral-8x7B-Instruct-v0.1
-          mistralai/Mixtral-8x7B-Instruct-v0.1
-        ]
-      end
-
-      # TODO(roman): Deprecated. Remove by Sept 2024
-      def self.available_ai_bots
-        %w[
-          gpt-3.5-turbo
-          gpt-4
-          gpt-4-turbo
-          gpt-4o
-          claude-2
-          gemini-1.5-pro
-          mixtral-8x7B-Instruct-V0.1
-          claude-3-opus
-          claude-3-sonnet
-          claude-3-haiku
-          cohere-command-r-plus
-        ]
-      end
     end
   end
 end