FEATURE: smarter persona tethering

SamSaffron · SamSaffron · commit 542f53c571b9 · 2024-10-14T18:06:04.000+11:00
Splits persona permissions so you can allow a persona on:

- chat dms
- personal messages
- topic mentions
- chat channels

(any combination is allowed)

Allow operator to tether a persona to an llm.


(Still TODO frontend changes)
diff --git a/app/models/ai_persona.rb b/app/models/ai_persona.rb
@@ -1,8 +1,8 @@
 # frozen_string_literal: true
 
 class AiPersona < ActiveRecord::Base
-  # TODO remove this line 01-11-2024
-  self.ignored_columns = [:commands]
+  # TODO remove this line 01-1-2025
+  self.ignored_columns = [:commands, :allow_chat, :mentionable]
 
   # places a hard limit, so per site we cache a maximum of 500 classes
   MAX_PERSONAS_PER_SITE = 500
@@ -52,49 +52,54 @@ def self.persona_users(user: nil)
       persona_cache[:persona_users] ||= AiPersona
         .where(enabled: true)
         .joins(:user)
-        .pluck(
-          "ai_personas.id, users.id, users.username_lower, allowed_group_ids, default_llm, mentionable, allow_chat",
-        )
-        .map do |id, user_id, username, allowed_group_ids, default_llm, mentionable, allow_chat|
+        .map do |persona|
           {
-            id: id,
-            user_id: user_id,
-            username: username,
-            allowed_group_ids: allowed_group_ids,
-            default_llm: default_llm,
-            mentionable: mentionable,
-            allow_chat: allow_chat,
+            id: persona.id,
+            user_id: persona.user_id,
+            username: persona.user.username_lower,
+            allowed_group_ids: persona.allowed_group_ids,
+            default_llm: persona.default_llm,
+            force_default_llm: persona.force_default_llm,
+            allow_chat_channel_mentions: persona.allow_chat_channel_mentions,
+            allow_chat_direct_messages: persona.allow_chat_direct_messages,
+            allow_topic_mentions: persona.allow_topic_mentions,
+            allow_personal_messages: persona.allow_personal_messages,
           }
         end
 
     if user
-      persona_users.select { |mentionable| user.in_any_groups?(mentionable[:allowed_group_ids]) }
+      persona_users.select { |user| user.in_any_groups?(mentionable[:allowed_group_ids]) }
     else
       persona_users
     end
   end
 
-  def self.allowed_chat(user: nil)
-    personas = persona_cache[:allowed_chat] ||= persona_users.select { |u| u[:allow_chat] }
+  def self.allowed_modalities(
+    user: nil,
+    allow_chat_channel_mentions: false,
+    allow_chat_direct_messages: false,
+    allow_topic_mentions: false,
+    allow_personal_messages: false
+  )
+    index =
+      "modality-#{allow_chat_channel_mentions}-#{allow_chat_direct_messages}-#{allow_topic_mentions}-#{allow_personal_messages}"
+
+    personas =
+      persona_cache[index.to_sym] ||= persona_users.select do |persona|
+        next true if allow_chat_channel_mentions && persona[:allow_chat_channel_mentions]
+        next true if allow_chat_direct_messages && persona[:allow_chat_direct_messages]
+        next true if allow_topic_mentions && persona[:allow_topic_mentions]
+        next true if allow_personal_messages && persona[:allow_personal_messages]
+        false
+      end
+
     if user
       personas.select { |u| user.in_any_groups?(u[:allowed_group_ids]) }
     else
       personas
     end
   end
 
-  def self.mentionables(user: nil)
-    all_mentionables =
-      persona_cache[:mentionables] ||= persona_users.select do |mentionable|
-        mentionable[:mentionable]
-      end
-    if user
-      all_mentionables.select { |mentionable| user.in_any_groups?(mentionable[:allowed_group_ids]) }
-    else
-      all_mentionables
-    end
-  end
-
   after_commit :bump_cache
 
   def bump_cache
@@ -113,7 +118,11 @@ def class_instance
       vision_max_pixels
       rag_conversation_chunks
       question_consolidator_llm
-      allow_chat
+      allow_chat_channel_mentions
+      allow_chat_direct_messages
+      allow_topic_mentions
+      allow_personal_messages
+      force_default_llm
       name
       description
       allowed_group_ids
@@ -243,7 +252,7 @@ def regenerate_rag_fragments
   private
 
   def chat_preconditions
-    if allow_chat && !default_llm
+    if (allow_chat_channel_mentions || allow_chat_direct_messages) && !default_llm
       errors.add(:default_llm, I18n.t("discourse_ai.ai_bot.personas.default_llm_required"))
     end
   end
@@ -281,7 +290,6 @@ def ensure_not_system
 #  temperature                 :float
 #  top_p                       :float
 #  user_id                     :integer
-#  mentionable                 :boolean          default(FALSE), not null
 #  default_llm                 :text
 #  max_context_posts           :integer
 #  max_post_context_tokens     :integer
@@ -291,16 +299,15 @@ def ensure_not_system
 #  rag_chunk_tokens            :integer          default(374), not null
 #  rag_chunk_overlap_tokens    :integer          default(10), not null
 #  rag_conversation_chunks     :integer          default(10), not null
-#  role                        :enum             default("bot"), not null
-#  role_category_ids           :integer          default([]), not null, is an Array
-#  role_tags                   :string           default([]), not null, is an Array
-#  role_group_ids              :integer          default([]), not null, is an Array
-#  role_whispers               :boolean          default(FALSE), not null
-#  role_max_responses_per_hour :integer          default(50), not null
 #  question_consolidator_llm   :text
-#  allow_chat                  :boolean          default(FALSE), not null
 #  tool_details                :boolean          default(TRUE), not null
 #  tools                       :json             not null
+#  forced_tool_count           :integer          default(-1), not null
+#  allow_chat_channel_mentions :boolean          default(FALSE), not null
+#  allow_chat_direct_messages  :boolean          default(FALSE), not null
+#  allow_topic_mentions        :boolean          default(FALSE), not null
+#  allow_personal_message      :boolean          default(TRUE), not null
+#  force_default_llm           :boolean          default(FALSE), not null
 #
 # Indexes
 #
diff --git a/db/migrate/20241014010245_ai_persona_chat_topic_refactor.rb b/db/migrate/20241014010245_ai_persona_chat_topic_refactor.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+class AiPersonaChatTopicRefactor < ActiveRecord::Migration[7.1]
+  def change
+    add_column :ai_personas, :allow_chat_channel_mentions, :boolean, default: false, null: false
+    add_column :ai_personas, :allow_chat_direct_messages, :boolean, default: false, null: false
+    add_column :ai_personas, :allow_topic_mentions, :boolean, default: false, null: false
+    add_column :ai_personas, :allow_personal_messages, :boolean, default: true, null: false
+    add_column :ai_personas, :force_default_llm, :boolean, default: false, null: false
+
+    execute <<~SQL
+      UPDATE ai_personas
+      SET allow_chat_channel_mentions = mentionable, allow_chat_direct_messages = true
+      WHERE allow_chat = true
+    SQL
+
+    execute <<~SQL
+      UPDATE ai_personas
+      SET allow_topic_mentions = true
+      WHERE mentionable = true
+    SQL
+  end
+end
diff --git a/db/post_migrate/20241014041242_ai_persona_post_migrate_drop_cols.rb b/db/post_migrate/20241014041242_ai_persona_post_migrate_drop_cols.rb
@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+class AiPersonaPostMigrateDropCols < ActiveRecord::Migration[7.1]
+  def change
+    remove_columns :ai_personas, :allow_chat
+    remove_columns :ai_personas, :mentionable
+  end
+end
diff --git a/lib/ai_bot/entry_point.rb b/lib/ai_bot/entry_point.rb
@@ -8,11 +8,7 @@ class EntryPoint
       Bot = Struct.new(:id, :name, :llm)
 
       def self.all_bot_ids
-        mentionable_persona_user_ids =
-          AiPersona.mentionables.map { |mentionable| mentionable[:user_id] }
-        mentionable_bot_users = LlmModel.joins(:user).pluck("users.id")
-
-        mentionable_bot_users + mentionable_persona_user_ids
+        AiPersona.persona_users.map { |persona| persona[:user_id] }
       end
 
       def self.find_participant_in(participant_ids)
@@ -140,7 +136,7 @@ def inject_into(plugin)
                 {
                   "id" => persona_user[:user_id],
                   "username" => persona_user[:username],
-                  "mentionable" => persona_user[:mentionable],
+                  "force_default_llm" => persona_user[:force_default_llm],
                   "is_persona" => true,
                 }
               end,
diff --git a/lib/ai_bot/personas/persona.rb b/lib/ai_bot/personas/persona.rb
@@ -21,7 +21,15 @@ def question_consolidator_llm
             nil
           end
 
-          def allow_chat
+          def force_default_llm
+            false
+          end
+
+          def allow_chat_channel_mentions
+            false
+          end
+
+          def allow_chat_direct_messages
             false
           end
 
diff --git a/lib/ai_bot/playground.rb b/lib/ai_bot/playground.rb
@@ -11,26 +11,34 @@ class Playground
 
       def self.find_chat_persona(message, channel, user)
         if channel.direct_message_channel?
-          AiPersona.allowed_chat.find do |p|
-            p[:user_id].in?(channel.allowed_user_ids) && (user.group_ids & p[:allowed_group_ids])
-          end
+          AiPersona
+            .allowed_modalities(allow_chat_direct_messages: true)
+            .find do |p|
+              p[:user_id].in?(channel.allowed_user_ids) && (user.group_ids & p[:allowed_group_ids])
+            end
         else
           # let's defer on the parse if there is no @ in the message
           if message.message.include?("@")
             mentions = message.parsed_mentions.parsed_direct_mentions
             if mentions.present?
-              AiPersona.allowed_chat.find do |p|
-                p[:username].in?(mentions) && (user.group_ids & p[:allowed_group_ids])
-              end
+              AiPersona
+                .allowed_modalities(allow_chat_channel_mentions: true)
+                .find { |p| p[:username].in?(mentions) && (user.group_ids & p[:allowed_group_ids]) }
             end
           end
         end
       end
 
       def self.schedule_chat_reply(message, channel, user, context)
         return if !SiteSetting.ai_bot_enabled
-        return if AiPersona.allowed_chat.blank?
-        return if AiPersona.allowed_chat.any? { |m| m[:user_id] == user.id }
+
+        all_chat =
+          AiPersona.allowed_modalities(
+            allow_chat_channel_mentions: true,
+            allow_chat_direct_messages: true,
+          )
+        return if all_chat.blank?
+        return if all_chat.any? { |m| m[:user_id] == user.id }
 
         persona = find_chat_persona(message, channel, user)
         return if !persona
@@ -56,15 +64,22 @@ def self.is_bot_user_id?(user_id)
 
       def self.schedule_reply(post)
         return if is_bot_user_id?(post.user_id)
+        mentionables = nil
 
-        bot_ids = LlmModel.joins(:user).pluck("users.id")
-        mentionables = AiPersona.mentionables(user: post.user)
+        if post.topic.private_message?
+          mentionables = AiPersona.allowed_modalities(user: post.user, allow_personal_messages: true)
+        else
+          mentionables = AiPersona.allowed_modalities(user: post.user, allow_topic_mentions: true)
+        end
 
         bot_user = nil
         mentioned = nil
 
+        all_llm_user_ids = LlmModel.joins(:user).pluck("users.id")
+
         if post.topic.private_message?
-          bot_user = post.topic.topic_allowed_users.where(user_id: bot_ids).first&.user
+          # this is an edge case, you started a PM with a different bot
+          bot_user = post.topic.topic_allowed_users.where(user_id: all_llm_user_ids).first&.user
           bot_user ||=
             post
               .topic
@@ -114,6 +129,10 @@ def self.schedule_reply(post)
 
           persona ||= DiscourseAi::AiBot::Personas::General
 
+          if persona && persona.force_default_llm
+            bot_user = User.find(persona.user_id)
+          end
+
           bot = DiscourseAi::AiBot::Bot.as(bot_user, persona: persona.new)
           new(bot).update_playground_with(post)
         end
diff --git a/spec/lib/modules/ai_bot/playground_spec.rb b/spec/lib/modules/ai_bot/playground_spec.rb
