FEATURE: Tool name validation (#842)

nvh0412 · web-flow · commit b60926c6e6a9 · 2025-02-07T14:34:47.000+11:00
* FEATURE: Tool name validation

- Add unique index to the name column of the ai_tools table
- correct our tests for AiToolController
- tool_name field which will be used to represent to LLM
- Add tool_name to Tools's presets
- Add duplicate tools validation for AiPersona
- Add unique constraint to the name column of the ai_tools table

* DEV: Validate duplicate tool_name between builin tools and custom tools

* lint

* chore: fix linting

* fix conlict mistakes

* chore: correct icon class

* chore: fix failed specs

* Add max_length to tool_name

* chore: correct the option name

* lintings

* fix lintings
diff --git a/app/controllers/discourse_ai/admin/ai_personas_controller.rb b/app/controllers/discourse_ai/admin/ai_personas_controller.rb
@@ -23,7 +23,12 @@ def index
           .each do |tool|
             tools << {
               id: "custom-#{tool.id}",
-              name: I18n.t("discourse_ai.tools.custom_name", name: tool.name.capitalize),
+              name:
+                I18n.t(
+                  "discourse_ai.tools.custom_name",
+                  name: tool.name.capitalize,
+                  tool_name: tool.tool_name,
+                ),
             }
           end
         llms =
diff --git a/app/controllers/discourse_ai/admin/ai_tools_controller.rb b/app/controllers/discourse_ai/admin/ai_tools_controller.rb
@@ -84,6 +84,7 @@ def ai_tool_params
           .require(:ai_tool)
           .permit(
             :name,
+            :tool_name,
             :description,
             :script,
             :summary,
diff --git a/app/models/ai_persona.rb b/app/models/ai_persona.rb
@@ -22,6 +22,9 @@ class AiPersona < ActiveRecord::Base
   validates :rag_chunk_overlap_tokens, numericality: { greater_than: -1, maximum: 200 }
   validates :rag_conversation_chunks, numericality: { greater_than: 0, maximum: 1000 }
   validates :forced_tool_count, numericality: { greater_than: -2, maximum: 100_000 }
+
+  validate :tools_can_not_be_duplicated
+
   has_many :rag_document_fragments, dependent: :destroy, as: :target
 
   belongs_to :created_by, class_name: "User"
@@ -107,6 +110,47 @@ def bump_cache
     self.class.persona_cache.flush!
   end
 
+  def tools_can_not_be_duplicated
+    return unless tools.is_a?(Array)
+
+    seen_tools = Set.new
+
+    custom_tool_ids = Set.new
+    builtin_tool_names = Set.new
+
+    tools.each do |tool|
+      inner_name, _, _ = tool.is_a?(Array) ? tool : [tool, nil]
+
+      if inner_name.start_with?("custom-")
+        custom_tool_ids.add(inner_name.split("-", 2).last.to_i)
+      else
+        builtin_tool_names.add(inner_name.downcase)
+      end
+
+      if seen_tools.include?(inner_name)
+        errors.add(:tools, I18n.t("discourse_ai.ai_bot.personas.cannot_have_duplicate_tools"))
+        break
+      else
+        seen_tools.add(inner_name)
+      end
+    end
+
+    return if errors.any?
+
+    # Checking if there are any duplicate tool_names between custom and builtin tools
+    if builtin_tool_names.present? && custom_tool_ids.present?
+      AiTool
+        .where(id: custom_tool_ids)
+        .pluck(:tool_name)
+        .each do |tool_name|
+          if builtin_tool_names.include?(tool_name.downcase)
+            errors.add(:tools, I18n.t("discourse_ai.ai_bot.personas.cannot_have_duplicate_tools"))
+            break
+          end
+        end
+    end
+  end
+
   def class_instance
     attributes = %i[
       id
diff --git a/app/models/ai_tool.rb b/app/models/ai_tool.rb
@@ -1,7 +1,8 @@
 # frozen_string_literal: true
 
 class AiTool < ActiveRecord::Base
-  validates :name, presence: true, length: { maximum: 100 }
+  validates :name, presence: true, length: { maximum: 100 }, uniqueness: true
+  validates :tool_name, presence: true, length: { maximum: 100 }
   validates :description, presence: true, length: { maximum: 1000 }
   validates :summary, presence: true, length: { maximum: 255 }
   validates :script, presence: true, length: { maximum: 100_000 }
@@ -12,8 +13,25 @@ class AiTool < ActiveRecord::Base
   has_many :uploads, through: :upload_references
   before_update :regenerate_rag_fragments
 
+  ALPHANUMERIC_PATTERN = /\A[a-zA-Z0-9_]+\z/
+
+  validates :tool_name,
+            format: {
+              with: ALPHANUMERIC_PATTERN,
+              message: I18n.t("discourse_ai.tools.name.characters"),
+            }
+
   def signature
-    { name: name, description: description, parameters: parameters.map(&:symbolize_keys) }
+    {
+      name: function_call_name,
+      description: description,
+      parameters: parameters.map(&:symbolize_keys),
+    }
+  end
+
+  # Backwards compatibility: if tool_name is not set (existing custom tools), use name
+  def function_call_name
+    tool_name.presence || name
   end
 
   def runner(parameters, llm:, bot_user:, context: {})
@@ -127,7 +145,8 @@ def self.presets
     [
       {
         preset_id: "browse_web_jina",
-        name: "browse_web",
+        name: "Browse Web",
+        tool_name: "browse_web",
         description: "Browse the web as a markdown document",
         parameters: [
           { name: "url", type: "string", required: true, description: "The URL to browse" },
@@ -148,7 +167,8 @@ def self.presets
       },
       {
         preset_id: "exchange_rate",
-        name: "exchange_rate",
+        name: "Exchange Rate",
+        tool_name: "exchange_rate",
         description: "Get current exchange rates for various currencies",
         parameters: [
           {
@@ -204,7 +224,8 @@ def self.presets
       },
       {
         preset_id: "stock_quote",
-        name: "stock_quote",
+        name: "Stock Quote (AlphaVantage)",
+        tool_name: "stock_quote",
         description: "Get real-time stock quote information using AlphaVantage API",
         parameters: [
           {
@@ -253,7 +274,8 @@ def self.presets
       },
       {
         preset_id: "image_generation",
-        name: "image_generation",
+        name: "Image Generation (Flux)",
+        tool_name: "image_generation",
         description:
           "Generate images using the FLUX model from Black Forest Labs using together.ai",
         parameters: [
@@ -348,4 +370,5 @@ def self.presets
 #  updated_at               :datetime         not null
 #  rag_chunk_tokens         :integer          default(374), not null
 #  rag_chunk_overlap_tokens :integer          default(10), not null
+#  tool_name                :string(100)      default(""), not null
 #
diff --git a/app/serializers/ai_custom_tool_serializer.rb b/app/serializers/ai_custom_tool_serializer.rb
@@ -3,6 +3,7 @@
 class AiCustomToolSerializer < ApplicationSerializer
   attributes :id,
              :name,
+             :tool_name,
              :description,
              :summary,
              :parameters,
diff --git a/assets/javascripts/discourse/admin/models/ai-tool.js b/assets/javascripts/discourse/admin/models/ai-tool.js
@@ -4,6 +4,7 @@ import RestModel from "discourse/models/rest";
 const CREATE_ATTRIBUTES = [
   "id",
   "name",
+  "tool_name",
   "description",
   "parameters",
   "script",
diff --git a/assets/javascripts/discourse/components/ai-tool-editor.gjs b/assets/javascripts/discourse/components/ai-tool-editor.gjs
@@ -83,6 +83,7 @@ export default class AiToolEditor extends Component {
     try {
       const data = this.editingModel.getProperties(
         "name",
+        "tool_name",
         "description",
         "parameters",
         "script",
@@ -178,6 +179,23 @@ export default class AiToolEditor extends Component {
           />
         </div>
 
+        <div class="control-group">
+          <label>{{i18n "discourse_ai.tools.tool_name"}}</label>
+          <input
+            {{on
+              "input"
+              (withEventValue (fn (mut this.editingModel.tool_name)))
+            }}
+            value={{this.editingModel.tool_name}}
+            type="text"
+            class="ai-tool-editor__tool_name"
+          />
+          <DTooltip
+            @icon="circle-question"
+            @content={{i18n "discourse_ai.tools.tool_name_help"}}
+          />
+        </div>
+
         <div class="control-group">
           <label>{{i18n "discourse_ai.tools.description"}}</label>
           <textarea
diff --git a/config/locales/client.en.yml b/config/locales/client.en.yml
@@ -292,11 +292,13 @@ en:
         short_title: "Tools"
         no_tools: "You have not created any tools yet"
         name: "Name"
-        subheader_description: "Tools extend the capabilities of AI bots with user-defined JavaScript functions."
+        name_help: "Name will show up in the Discourse UI and is the short identifier you will use to find the tool in various settings, it should be distinct (it is required)"
         new: "New tool"
-        name_help: "The unique name of the tool as used by the language model"
+        tool_name: "Tool Name"
+        tool_name_help: "Tool Name is presented to the large language model. It is not distinct, but it is distinct per persona. (persona validates on save)"
         description: "Description"
         description_help: "A clear description of the tool's purpose for the language model"
+        subheader_description: "Tools extend the capabilities of AI bots with user-defined JavaScript functions."
         summary: "Summary"
         summary_help: "Summary of tools purpose to be displayed to end users"
         script: "Script"
diff --git a/config/locales/server.en.yml b/config/locales/server.en.yml
@@ -215,6 +215,8 @@ en:
           name: "Flux image generator (Together.ai)"
         empty_tool:
           name: "Start from blank..."
+      name:
+        characters: "must only include numbers, letters, and underscores"
 
     ai_helper:
       errors:
@@ -260,6 +262,7 @@ en:
         default_llm_required: "Default LLM model is required prior to enabling Chat"
         cannot_delete_system_persona: "System personas cannot be deleted, please disable it instead"
         cannot_edit_system_persona: "System personas can only be renamed, you may not edit tools or system prompt, instead disable and make a copy"
+        cannot_have_duplicate_tools: "Can not have duplicate tools"
         github_helper:
           name: "GitHub Helper"
           description: "AI Bot specialized in assisting with GitHub-related tasks and questions"
diff --git a/db/migrate/20241020010245_add_tool_name_to_ai_tools.rb b/db/migrate/20241020010245_add_tool_name_to_ai_tools.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+class AddToolNameToAiTools < ActiveRecord::Migration[7.1]
+  def up
+    add_column :ai_tools,
+               :tool_name,
+               :string,
+               null: false,
+               limit: 100,
+               default: "",
+               if_not_exists: true
+
+    # Migrate existing name to tool_name
+    execute <<~SQL
+      UPDATE ai_tools
+      SET tool_name = regexp_replace(LOWER(name),'[^a-z0-9_]','', 'g');
+    SQL
+  end
+
+  def down
+    remove_column :ai_tools, :tool_name, if_exists: true
+  end
+end
diff --git a/db/migrate/20241023041242_add_unique_constraint_to_ai_tools.rb b/db/migrate/20241023041242_add_unique_constraint_to_ai_tools.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+class AddUniqueConstraintToAiTools < ActiveRecord::Migration[7.1]
+  def up
+    # We need to remove duplicates before adding the unique constraint
+    execute <<~SQL
+      WITH duplicates AS (
+        SELECT name, COUNT(*) as count, MIN(id) as keeper_id
+        FROM ai_tools
+        GROUP BY name
+        HAVING COUNT(*) > 1
+      )
+      UPDATE ai_tools AS p
+      SET name = CONCAT(p.name, p.id)
+      FROM duplicates d
+      WHERE p.name = d.name
+      AND p.id != d.keeper_id;
+    SQL
+
+    add_index :ai_personas, :name, unique: true, if_not_exists: true
+  end
+
+  def down
+    remove_index :ai_personas, :name, if_exists: true
+  end
+end
diff --git a/lib/ai_bot/personas/general.rb b/lib/ai_bot/personas/general.rb
@@ -10,7 +10,6 @@ def tools
             Tools::Google,
             Tools::Image,
             Tools::Read,
-            Tools::Image,
             Tools::ListCategories,
             Tools::ListTags,
           ]
diff --git a/lib/ai_bot/tools/custom.rb b/lib/ai_bot/tools/custom.rb
@@ -26,8 +26,11 @@ def self.signature
           AiTool.find(tool_id).signature
         end
 
+        # Backwards compatibility: if tool_name is not set (existing custom tools), use name
         def self.name
-          AiTool.where(id: tool_id).pluck(:name).first
+          name, tool_name = AiTool.where(id: tool_id).pluck(:name, :tool_name).first
+
+          tool_name.presence || name
         end
 
         def initialize(*args, **kwargs)
diff --git a/spec/fabricators/ai_tool_fabricator.rb b/spec/fabricators/ai_tool_fabricator.rb
@@ -0,0 +1,10 @@
+# frozen_string_literal: true
+
+Fabricator(:ai_tool) do
+  name "github tool"
+  tool_name "github_tool"
+  description "This is a tool for GitHub"
+  summary "This is a tool for GitHub"
+  script "puts 'Hello, GitHub!'"
+  created_by_id 1
+end
diff --git a/spec/lib/modules/ai_bot/playground_spec.rb b/spec/lib/modules/ai_bot/playground_spec.rb
@@ -82,6 +82,7 @@
     let!(:custom_tool) do
       AiTool.create!(
         name: "search",
+        tool_name: "search",
         summary: "searching for things",
         description: "A test custom tool",
         parameters: [{ name: "query", type: "string", description: "Input for the custom tool" }],
diff --git a/spec/models/ai_persona_spec.rb b/spec/models/ai_persona_spec.rb
@@ -24,6 +24,53 @@
     expect(persona.valid?).to eq(true)
   end
 
+  it "validates tools" do
+    persona =
+      AiPersona.new(
+        name: "test",
+        description: "test",
+        system_prompt: "test",
+        tools: [],
+        allowed_group_ids: [],
+      )
+
+    Fabricate(:ai_tool, id: 1)
+    Fabricate(:ai_tool, id: 2, name: "Archie search", tool_name: "search")
+
+    expect(persona.valid?).to eq(true)
+
+    persona.tools = %w[search image_generation]
+    expect(persona.valid?).to eq(true)
+
+    persona.tools = %w[search image_generation search]
+    expect(persona.valid?).to eq(false)
+    expect(persona.errors[:tools]).to eq(["Can not have duplicate tools"])
+
+    persona.tools = [["custom-1", { test: "test" }, false], ["custom-2", { test: "test" }, false]]
+    expect(persona.valid?).to eq(true)
+    expect(persona.errors[:tools]).to eq([])
+
+    persona.tools = [["custom-1", { test: "test" }, false], ["custom-1", { test: "test" }, false]]
+    expect(persona.valid?).to eq(false)
+    expect(persona.errors[:tools]).to eq(["Can not have duplicate tools"])
+
+    persona.tools = [
+      ["custom-1", { test: "test" }, false],
+      ["custom-2", { test: "test" }, false],
+      "image_generation",
+    ]
+    expect(persona.valid?).to eq(true)
+    expect(persona.errors[:tools]).to eq([])
+
+    persona.tools = [
+      ["custom-1", { test: "test" }, false],
+      ["custom-2", { test: "test" }, false],
+      "Search",
+    ]
+    expect(persona.valid?).to eq(false)
+    expect(persona.errors[:tools]).to eq(["Can not have duplicate tools"])
+  end
+
   it "allows creation of user" do
     persona =
       AiPersona.create!(
diff --git a/spec/models/ai_tool_spec.rb b/spec/models/ai_tool_spec.rb
diff --git a/spec/requests/admin/ai_tools_controller_spec.rb b/spec/requests/admin/ai_tools_controller_spec.rb
