DRAFT: Create AI Bot users dynamically and support custom LlmModels

Author: romanrizzi

app/controllers/discourse_ai/ai_bot/bot_controller.rb

@@ -31,12 +31,10 @@ def stop_streaming_response
       end
 
       def show_bot_username
-        bot_user_id = DiscourseAi::AiBot::EntryPoint.map_bot_model_to_user_id(params[:username])
-        raise Discourse::InvalidParameters.new(:username) if !bot_user_id
+        bot_user = DiscourseAi::AiBot::EntryPoint.find_user_from_model(params[:username])
+        raise Discourse::InvalidParameters.new(:username) if !bot_user
 
-        bot_username_lower = User.find(bot_user_id).username_lower
-
-        render json: { bot_username: bot_username_lower }, status: 200
+        render json: { bot_username: bot_user.username_lower }, status: 200
       end
     end
   end

app/models/shared_ai_conversation.rb

@@ -133,12 +133,10 @@ def formatted_excerpt
   end
 
   def self.build_conversation_data(topic, max_posts: DEFAULT_MAX_POSTS, include_usernames: false)
-    llm_name = nil
-    topic.topic_allowed_users.each do |tu|
-      if DiscourseAi::AiBot::EntryPoint::BOT_USER_IDS.include?(tu.user_id)
-        llm_name = DiscourseAi::AiBot::EntryPoint.find_bot_by_id(tu.user_id)&.llm
-      end
-    end
+    allowed_user_ids = topic.topic_allowed_users.pluck(:user_id)
+    ai_bot_participant = DiscourseAi::AiBot::EntryPoint.find_participant_in(allowed_user_ids)
+
+    llm_name = ai_bot_participant&.llm
 
     llm_name = ActiveSupport::Inflector.humanize(llm_name) if llm_name
     llm_name ||= I18n.t("discourse_ai.unknown_model")
@@ -170,9 +168,7 @@ def self.build_conversation_data(topic, max_posts: DEFAULT_MAX_POSTS, include_us
             cooked: post.cooked,
           }
 
-          mapped[:persona] = persona if ::DiscourseAi::AiBot::EntryPoint::BOT_USER_IDS.include?(
-            post.user_id,
-          )
+          mapped[:persona] = persona if ai_bot_participant&.id == post.user_id
           mapped[:username] = post.user&.username if include_usernames
           mapped
         end,

config/settings.yml

@@ -354,18 +354,7 @@ discourse_ai:
     type: list
     default: "gpt-3.5-turbo"
     client: true
-    choices:
-      - 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
+    choices: "DiscourseAi::Configuration::LlmEnumerator.ai_bot_models"
   ai_bot_add_to_header:
     default: true
     client: true

db/post_migrate/20240523142424_custom_fields_to_track_ai_bot_users.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+class CustomFieldsToTrackAiBotUsers < ActiveRecord::Migration[7.0]
+  def up
+    existing_bot_user_ids = DB.query_single("SELECT id FROM users WHERE id <= -110 AND id >= -121")
+
+    custom_field_rows =
+      existing_bot_user_ids
+        .map { |id| "(bot_model_name, #{id_to_model_name(id)}, #{id})" }
+        .join(",")
+
+    DB.exec(<<~SQL, rows: custom_field_rows) if custom_field_rows.present?
+        INSERT INTO user_custom_fields (name, value, user_id)
+        VALUES :rows;
+      SQL
+  end
+
+  def id_to_model_name(id)
+    # Skip -116. fake model.
+    case id
+    when -110
+      "gpt-4"
+    when -111
+      "gpt-3.5-turbo"
+    when -112
+      "claude-2"
+    when -113
+      "gpt-4-turbo"
+    when -114
+      "mixtral-8x7B-Instruct-V0.1"
+    when -115
+      "gemini-1.5-pro"
+    when -116
+      "fake"
+    when -117
+      "claude-3-opus"
+    when -118
+      "claude-3-sonnet"
+    when -119
+      "claude-3-haiku"
+    when -120
+      "cohere-command-r-plus"
+    else
+      "gpt-4o"
+    end
+  end
+
+  def down
+    raise ActiveRecord::IrreversibleMigration
+  end
+end

lib/ai_bot/bot.rb

@@ -163,76 +163,55 @@ def invoke_tool(tool, llm, cancel, context, &update_blk)
       end
 
       def self.guess_model(bot_user)
-        # HACK(roman): We'll do this until we define how we represent different providers in the bot settings
-        guess =
-          case bot_user.id
-          when DiscourseAi::AiBot::EntryPoint::CLAUDE_V2_ID
-            if DiscourseAi::Completions::Endpoints::AwsBedrock.correctly_configured?("claude-2")
-              "aws_bedrock:claude-2"
-            else
-              "anthropic:claude-2"
-            end
-          when DiscourseAi::AiBot::EntryPoint::GPT4_ID
-            "open_ai:gpt-4"
-          when DiscourseAi::AiBot::EntryPoint::GPT4_TURBO_ID
-            "open_ai:gpt-4-turbo"
-          when DiscourseAi::AiBot::EntryPoint::GPT4O_ID
-            "open_ai:gpt-4o"
-          when DiscourseAi::AiBot::EntryPoint::GPT3_5_TURBO_ID
-            "open_ai:gpt-3.5-turbo-16k"
-          when DiscourseAi::AiBot::EntryPoint::MIXTRAL_ID
-            mixtral_model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
+        base_choice = guess_base_model(bot_user)
+
+        if base_choice
+          provider, model_name = base_choice.split(":")
+          llm_model = LlmModel.find_by(provider: provider, name: model_name)
+
+          return "custom:#{llm_model.id}" if llm_model
+        end
+
+        base_choice
+      end
+
+      # HACK(roman): We'll do this until we define how we represent different providers in the bot settings
+      def self.guess_base_model(bot_user)
+        associated_llm =
+          bot_user.custom_fields[DiscourseAi::AiBot::EntryPoint::BOT_MODEL_CUSTOM_FIELD]
+
+        return if associated_llm.nil? # Might be a persona user. Handled by constructor.
+
+        return "open_ai:gpt-3.5-turbo-16k" if associated_llm == "gpt-3.5-turbo"
+        return "open_ai:#{associated_llm}" if associated_llm.starts_with?("gpt-")
+
+        return "google:gemini-1.5-pro" if associated_llm == "gemini-1.5-pro"
+
+        return "fake:fake" if associated_llm == "fake"
+
+        return "cohere:command-r-plus" if associated_llm = "command-r-plus"
+
+        if %w[claude-3-opus claude-3-sonnet claude-3-haiku].include?(associated_llm)
+          if DiscourseAi::Completions::Endpoints::AwsBedrock.correctly_configured?(associated_llm)
+            return "aws_bedrock:#{associated_llm}"
+          else
+            return "anthropic:#{associated_llm}"
+          end
+        end
+
+        if associated_llm == "mistralai/Mixtral-8x7B-Instruct-v0.1"
+          if DiscourseAi::Completions::Endpoints::Vllm.correctly_configured?(associated_llm)
             if DiscourseAi::Completions::Endpoints::Vllm.correctly_configured?(mixtral_model)
-              "vllm:#{mixtral_model}"
+              return "vllm:#{associated_llm}"
             elsif DiscourseAi::Completions::Endpoints::HuggingFace.correctly_configured?(
-                  mixtral_model,
+                  associated_llm,
                 )
-              "hugging_face:#{mixtral_model}"
-            else
-              "ollama:mistral"
-            end
-          when DiscourseAi::AiBot::EntryPoint::GEMINI_ID
-            "google:gemini-1.5-pro"
-          when DiscourseAi::AiBot::EntryPoint::FAKE_ID
-            "fake:fake"
-          when DiscourseAi::AiBot::EntryPoint::CLAUDE_3_OPUS_ID
-            if DiscourseAi::Completions::Endpoints::AwsBedrock.correctly_configured?(
-                 "claude-3-opus",
-               )
-              "aws_bedrock:claude-3-opus"
-            else
-              "anthropic:claude-3-opus"
-            end
-          when DiscourseAi::AiBot::EntryPoint::COHERE_COMMAND_R_PLUS
-            "cohere:command-r-plus"
-          when DiscourseAi::AiBot::EntryPoint::CLAUDE_3_SONNET_ID
-            if DiscourseAi::Completions::Endpoints::AwsBedrock.correctly_configured?(
-                 "claude-3-sonnet",
-               )
-              "aws_bedrock:claude-3-sonnet"
+              return "hugging_face:#{associated_llm}"
             else
-              "anthropic:claude-3-sonnet"
+              return "ollama:mistral"
             end
-          when DiscourseAi::AiBot::EntryPoint::CLAUDE_3_HAIKU_ID
-            if DiscourseAi::Completions::Endpoints::AwsBedrock.correctly_configured?(
-                 "claude-3-haiku",
-               )
-              "aws_bedrock:claude-3-haiku"
-            else
-              "anthropic:claude-3-haiku"
-            end
-          else
-            nil
           end
-
-        if guess
-          provider, model_name = guess.split(":")
-          llm_model = LlmModel.find_by(provider: provider, name: model_name)
-
-          return "custom:#{llm_model.id}" if llm_model
         end
-
-        guess
       end
 
       def build_placeholder(summary, details, custom_raw: nil)

lib/ai_bot/entry_point.rb

@@ -6,82 +6,58 @@ module AiBot
 
     class EntryPoint
       REQUIRE_TITLE_UPDATE = "discourse-ai-title-update"
-
-      GPT4_ID = -110
-      GPT3_5_TURBO_ID = -111
-      CLAUDE_V2_ID = -112
-      GPT4_TURBO_ID = -113
-      MIXTRAL_ID = -114
-      GEMINI_ID = -115
-      FAKE_ID = -116 # only used for dev and test
-      CLAUDE_3_OPUS_ID = -117
-      CLAUDE_3_SONNET_ID = -118
-      CLAUDE_3_HAIKU_ID = -119
-      COHERE_COMMAND_R_PLUS = -120
-      GPT4O_ID = -121
-
-      BOTS = [
-        [GPT4_ID, "gpt4_bot", "gpt-4"],
-        [GPT3_5_TURBO_ID, "gpt3.5_bot", "gpt-3.5-turbo"],
-        [CLAUDE_V2_ID, "claude_bot", "claude-2"],
-        [GPT4_TURBO_ID, "gpt4t_bot", "gpt-4-turbo"],
-        [MIXTRAL_ID, "mixtral_bot", "mixtral-8x7B-Instruct-V0.1"],
-        [GEMINI_ID, "gemini_bot", "gemini-1.5-pro"],
-        [FAKE_ID, "fake_bot", "fake"],
-        [CLAUDE_3_OPUS_ID, "claude_3_opus_bot", "claude-3-opus"],
-        [CLAUDE_3_SONNET_ID, "claude_3_sonnet_bot", "claude-3-sonnet"],
-        [CLAUDE_3_HAIKU_ID, "claude_3_haiku_bot", "claude-3-haiku"],
-        [COHERE_COMMAND_R_PLUS, "cohere_command_bot", "cohere-command-r-plus"],
-        [GPT4O_ID, "gpt4o_bot", "gpt-4o"],
-      ]
-
-      BOT_USER_IDS = BOTS.map(&:first)
-
+      BOT_MODEL_CUSTOM_FIELD = "bot_model_name"
       Bot = Struct.new(:id, :name, :llm)
 
       def self.all_bot_ids
-        BOT_USER_IDS.concat(AiPersona.mentionables.map { |mentionable| mentionable[:user_id] })
+        mentionable_persona_user_ids =
+          AiPersona.mentionables.map { |mentionable| mentionable[:user_id] }
+        mentionable_bot_users =
+          User
+            .joins(:_custom_fields)
+            .where(active: true, user_custom_fields: { name: BOT_MODEL_CUSTOM_FIELD })
+            .pluck(:user_id)
+
+        mentionable_bot_users + mentionable_persona_user_ids
       end
 
-      def self.find_bot_by_id(id)
-        found = DiscourseAi::AiBot::EntryPoint::BOTS.find { |bot| bot[0] == id }
-        return if !found
-        Bot.new(found[0], found[1], found[2])
+      def self.find_participant_in(participant_ids)
+        participant_data =
+          UserCustomField
+            .includes(:user)
+            .where(users: { active: true }, name: BOT_MODEL_CUSTOM_FIELD, user_id: participant_ids)
+            .last
+
+        return if participant_data.nil?
+
+        Bot.new(
+          participant_data.user.id,
+          participant_data.user.username_lower,
+          participant_data.value,
+        )
       end
 
-      def self.map_bot_model_to_user_id(model_name)
-        case model_name
-        in "gpt-4o"
-          GPT4O_ID
-        in "gpt-4-turbo"
-          GPT4_TURBO_ID
-        in "gpt-3.5-turbo"
-          GPT3_5_TURBO_ID
-        in "gpt-4"
-          GPT4_ID
-        in "claude-2"
-          CLAUDE_V2_ID
-        in "mixtral-8x7B-Instruct-V0.1"
-          MIXTRAL_ID
-        in "gemini-1.5-pro"
-          GEMINI_ID
-        in "fake"
-          FAKE_ID
-        in "claude-3-opus"
-          CLAUDE_3_OPUS_ID
-        in "claude-3-sonnet"
-          CLAUDE_3_SONNET_ID
-        in "claude-3-haiku"
-          CLAUDE_3_HAIKU_ID
-        in "cohere-command-r-plus"
-          COHERE_COMMAND_R_PLUS
-        else
-          nil
-        end
+      def self.find_user_from_model(model_name)
+        UserCustomField
+          .includes(:user)
+          .find_by(name: BOT_MODEL_CUSTOM_FIELD, value: model_name)
+          &.user
+      end
+
+      def self.enabled_user_ids_and_models_map
+        enabled_models = SiteSetting.ai_bot_enabled_chat_bots.split("|")
+
+        DB.query_hash(<<~SQL, model_names: enabled_models, cf_name: BOT_MODEL_CUSTOM_FIELD)
+          SELECT users.username AS username, users.id AS id, ucf.value AS model_name
+          FROM user_custom_fields ucf
+          INNER JOIN users ON ucf.user_id = users.id
+          WHERE ucf.name = :cf_name 
+          AND ucf.value IN (:model_names)
+        SQL
       end
 
       # Most errors are simply "not_allowed"
-      # we do not want to reveal information about this sytem
+      # we do not want to reveal information about this system
       # the 2 exceptions are "other_people_in_pm" and "other_content_in_pm"
       # in both cases you have access to the PM so we are not revealing anything
       def self.ai_share_error(topic, guardian)
@@ -170,25 +146,11 @@ def inject_into(plugin)
               scope.user.in_any_groups?(SiteSetting.ai_bot_allowed_groups_map)
           end,
         ) do
-          model_map = {}
-          SiteSetting
-            .ai_bot_enabled_chat_bots
-            .split("|")
-            .each do |bot_name|
-              model_map[
-                ::DiscourseAi::AiBot::EntryPoint.map_bot_model_to_user_id(bot_name)
-              ] = bot_name
-            end
-
-          # not 100% ideal, cause it is one extra query, but we need it
-          bots = DB.query_hash(<<~SQL, user_ids: model_map.keys)
-            SELECT username, id FROM users WHERE id IN (:user_ids)
-          SQL
+          bots_map = ::DiscourseAi::AiBot::EntryPoint.enabled_user_ids_and_models_map
 
-          bots.each { |hash| hash["model_name"] = model_map[hash["id"]] }
           persona_users = AiPersona.persona_users(user: scope.user)
           if persona_users.present?
-            bots.concat(
+            bots_map.concat(
               persona_users.map do |persona_user|
                 {
                   "id" => persona_user[:user_id],
@@ -198,7 +160,8 @@ def inject_into(plugin)
               end,
             )
           end
-          bots
+
+          bots_map
         end
 
         plugin.add_to_serializer(:current_user, :can_use_assistant) do