FEATURE: correctly decorate AI bots

assets/javascripts/discourse/lib/ai-bot-helper.js

@@ -7,21 +7,40 @@ import ShareFullTopicModal from "../components/modal/share-full-topic-modal";
 
 const MAX_PERSONA_USER_ID = -1200;
 
-let enabledChatBotIds;
+let enabledChatBotMap = null;
+
+function ensureBotMap() {
+  if (!enabledChatBotMap) {
+    const currentUser = getOwnerWithFallback(this).lookup(
+      "service:current-user"
+    );
+    enabledChatBotMap = {};
+    currentUser.ai_enabled_chat_bots.forEach((bot) => {
+      enabledChatBotMap[bot.id] = bot;
+    });
+  }
+}
 
 export function isGPTBot(user) {
   if (!user) {
     return;
   }
 
-  if (!enabledChatBotIds) {
-    const currentUser = getOwnerWithFallback(this).lookup(
-      "service:current-user"
-    );
-    enabledChatBotIds = currentUser.ai_enabled_chat_bots.map((bot) => bot.id);
+  ensureBotMap();
+  return !!enabledChatBotMap[user.id];
+}
+
+export function getBotType(user) {
+  if (!user) {
+    return;
   }
 
-  return enabledChatBotIds.includes(user.id);
+  ensureBotMap();
+  const bot = enabledChatBotMap[user.id];
+  if (!bot) {
+    return;
+  }
+  return bot.is_persona ? "persona" : "llm";
 }
 
 export function isPostFromAiBot(post, currentUser) {

assets/javascripts/initializers/ai-bot-replies.js

@@ -8,6 +8,7 @@ import AiCancelStreamingButton from "../discourse/components/post-menu/ai-cancel
 import AiDebugButton from "../discourse/components/post-menu/ai-debug-button";
 import AiShareButton from "../discourse/components/post-menu/ai-share-button";
 import {
+  getBotType,
   isGPTBot,
   showShareConversationModal,
 } from "../discourse/lib/ai-bot-helper";
@@ -62,13 +63,19 @@ function initializePersonaDecorator(api) {
 
 function initializeWidgetPersonaDecorator(api) {
   api.decorateWidget(`poster-name:after`, (dec) => {
-    if (!isGPTBot(dec.attrs.user)) {
-      return;
+    const botType = getBotType(dec.attrs.user);
+    // we have 2 ways of decorating
+    // 1. if a bot is a LLM we decorate with persona name
+    // 2. if bot is a persona we decorate with LLM name
+    if (botType === "llm") {
+      return dec.widget.attach("persona-flair", {
+        personaName: dec.model?.topic?.ai_persona_name,
+      });
+    } else if (botType === "persona") {
+      return dec.widget.attach("persona-flair", {
+        personaName: dec.model?.llm_name,
+      });
     }
-
-    return dec.widget.attach("persona-flair", {
-      personaName: dec.model?.topic?.ai_persona_name,
-    });
   });
 
   registerWidgetShim(

assets/stylesheets/modules/ai-bot/common/bot-replies.scss

@@ -78,7 +78,6 @@ article.streaming nav.post-controls .actions button.cancel-streaming {
 .topic-body .persona-flair {
   order: 2;
   font-size: var(--font-down-1);
-  padding-top: 3px;
 }
 
 details.ai-quote {

config/locales/server.en.yml

@@ -113,6 +113,7 @@ en:
     ai_discord_search_mode: "Select the search mode to use for Discord search"
     ai_discord_search_persona: "The persona to use for Discord search."
     ai_discord_allowed_guilds: "Discord guilds (servers) where the bot is allowed to search"
+    ai_enable_experimental_bot_ux: "Enable experimental bot UI that allows for a more dedicated experience"
 
   reviewables:
     reasons:

lib/ai_bot/entry_point.rb

@@ -4,6 +4,7 @@ module DiscourseAi
   module AiBot
     USER_AGENT = "Discourse AI Bot 1.0 (https://www.discourse.org)"
     TOPIC_AI_BOT_PM_FIELD = "is_ai_bot_pm"
+    POST_AI_LLM_NAME_FIELD = "ai_llm_name"
 
     class EntryPoint
       Bot = Struct.new(:id, :name, :llm)
@@ -65,6 +66,10 @@ def self.ai_share_error(topic, guardian)
       end
 
       def inject_into(plugin)
+        # Long term we need a better API here
+        # we only want to load this custom field for bots
+        TopicView.default_post_custom_fields << POST_AI_LLM_NAME_FIELD
+
         plugin.register_topic_custom_field_type(TOPIC_AI_BOT_PM_FIELD, :string)
 
         plugin.on(:topic_created) do |topic|
@@ -139,6 +144,14 @@ def inject_into(plugin)
           end,
         ) { true }
 
+        plugin.add_to_serializer(
+          :post,
+          :llm_name,
+          include_condition: -> do
+            object.topic.private_message? && object.custom_fields[POST_AI_LLM_NAME_FIELD]
+          end,
+        ) { object.custom_fields[POST_AI_LLM_NAME_FIELD] }
+
         plugin.add_to_serializer(
           :current_user,
           :ai_enabled_personas,

lib/ai_bot/playground.rb

@@ -458,6 +458,9 @@ def reply_to(
               skip_jobs: true,
               post_type: post_type,
               skip_guardian: true,
+              custom_fields: {
+                DiscourseAi::AiBot::POST_AI_LLM_NAME_FIELD => bot.llm.llm_model.name,
+              },
             )
 
           publish_update(reply_post, { raw: reply_post.cooked })

spec/lib/modules/ai_bot/playground_spec.rb

@@ -783,6 +783,10 @@
       last_post = post.topic.posts.order(:post_number).last
       expect(last_post.raw).to eq("Yes I can")
       expect(last_post.user_id).to eq(persona.user_id)
+
+      expect(last_post.custom_fields[DiscourseAi::AiBot::POST_AI_LLM_NAME_FIELD]).to eq(
+        gpt_35_turbo.name,
+      )
     end
 
     it "picks the correct llm for persona in PMs" do

spec/requests/ai_bot/topic_serialization_spec.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+RSpec.describe "AI Bot Post Serializer" do
+  fab!(:current_user) { Fabricate(:user) }
+  fab!(:bot_user) { Fabricate(:user) }
+
+  before do
+    SiteSetting.ai_bot_enabled = true
+    sign_in(current_user)
+  end
+
+  describe "llm_name in post serializer" do
+    it "includes llm_name when custom field is set in a PM" do
+      pm_topic = Fabricate(:private_message_topic, user: current_user)
+
+      # Create a bot post with the custom field set
+      bot_post =
+        Fabricate(
+          :post,
+          topic: pm_topic,
+          user: bot_user,
+          custom_fields: {
+            DiscourseAi::AiBot::POST_AI_LLM_NAME_FIELD => "bob",
+          },
+        )
+
+      get "/t/#{pm_topic.id}.json"
+      expect(response.status).to eq(200)
+
+      json = response.parsed_body
+      bot_post_data = json["post_stream"]["posts"].find { |p| p["id"] == bot_post.id }
+
+      expect(bot_post_data).to have_key("llm_name")
+      expect(bot_post_data["llm_name"]).to eq("bob")
+    end
+  end
+end