From b60926c6e6a91c54220112ed02e7a37229401165 Mon Sep 17 00:00:00 2001
From: Hoa Nguyen <nvh0412@gmail.com>
Date: Fri, 7 Feb 2025 14:34:47 +1100
Subject: [PATCH] FEATURE: Tool name validation (#842)

* 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
---
 .../admin/ai_personas_controller.rb           |  7 ++-
 .../discourse_ai/admin/ai_tools_controller.rb |  1 +
 app/models/ai_persona.rb                      | 44 +++++++++++++++++
 app/models/ai_tool.rb                         | 35 +++++++++++---
 app/serializers/ai_custom_tool_serializer.rb  |  1 +
 .../discourse/admin/models/ai-tool.js         |  1 +
 .../discourse/components/ai-tool-editor.gjs   | 18 +++++++
 config/locales/client.en.yml                  |  6 ++-
 config/locales/server.en.yml                  |  3 ++
 ...0241020010245_add_tool_name_to_ai_tools.rb | 23 +++++++++
 ...41242_add_unique_constraint_to_ai_tools.rb | 25 ++++++++++
 lib/ai_bot/personas/general.rb                |  1 -
 lib/ai_bot/tools/custom.rb                    |  5 +-
 spec/fabricators/ai_tool_fabricator.rb        | 10 ++++
 spec/lib/modules/ai_bot/playground_spec.rb    |  1 +
 spec/models/ai_persona_spec.rb                | 47 +++++++++++++++++++
 spec/models/ai_tool_spec.rb                   |  5 +-
 .../admin/ai_tools_controller_spec.rb         |  7 ++-
 18 files changed, 225 insertions(+), 15 deletions(-)
 create mode 100644 db/migrate/20241020010245_add_tool_name_to_ai_tools.rb
 create mode 100644 db/migrate/20241023041242_add_unique_constraint_to_ai_tools.rb
 create mode 100644 spec/fabricators/ai_tool_fabricator.rb

diff --git a/app/controllers/discourse_ai/admin/ai_personas_controller.rb b/app/controllers/discourse_ai/admin/ai_personas_controller.rb
index 2d07b7bbe..b0317f02f 100644
--- 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
index c8a85c60e..fc7b29844 100644
--- 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
index cc1df7e02..2c573977e 100644
--- 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
index ee2026cf9..97b2a983f 100644
--- 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
index 8fe1fd412..1ac804300 100644
--- 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
index bfeea21bf..7188a68c1 100644
--- 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
index ef7c16f37..154ce390b 100644
--- 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
index fb043d832..3f2f776bd 100644
--- 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
index a00b56e27..05ad8eed4 100644
--- 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
new file mode 100644
index 000000000..ab9a549fc
--- /dev/null
+++ 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
new file mode 100644
index 000000000..2c1a37a07
--- /dev/null
+++ 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
index dfb79ad6c..3e06a3f72 100644
--- 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
index 5c77860f4..8a4fbc999 100644
--- 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
new file mode 100644
index 000000000..dc464fe77
--- /dev/null
+++ 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
index b04d36553..4a3e541cf 100644
--- 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
index 38ab93fcb..2a7e9921f 100644
--- 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
index fda5351fc..1d7135cbd 100644
--- a/spec/models/ai_tool_spec.rb
+++ b/spec/models/ai_tool_spec.rb
@@ -11,7 +11,8 @@ def create_tool(
     rag_chunk_overlap_tokens: nil
   )
     AiTool.create!(
-      name: "test",
+      name: "test #{SecureRandom.uuid}",
+      tool_name: "test_#{SecureRandom.uuid.underscore}",
       description: "test",
       parameters: parameters || [{ name: "query", type: "string", desciption: "perform a search" }],
       script: script || "function invoke(params) { return params; }",
@@ -27,7 +28,7 @@ def create_tool(
 
     expect(tool.signature).to eq(
       {
-        name: "test",
+        name: tool.tool_name,
         description: "test",
         parameters: [{ name: "query", type: "string", desciption: "perform a search" }],
       },
diff --git a/spec/requests/admin/ai_tools_controller_spec.rb b/spec/requests/admin/ai_tools_controller_spec.rb
index 6a8ce92e9..ef9d87370 100644
--- a/spec/requests/admin/ai_tools_controller_spec.rb
+++ b/spec/requests/admin/ai_tools_controller_spec.rb
@@ -5,6 +5,7 @@
   fab!(:ai_tool) do
     AiTool.create!(
       name: "Test Tool",
+      tool_name: "test_tool",
       description: "A test tool",
       script: "function invoke(params) { return params; }",
       parameters: [
@@ -46,7 +47,8 @@
   describe "POST #create" do
     let(:valid_attributes) do
       {
-        name: "Test Tool",
+        name: "Test Tool 1",
+        tool_name: "test_tool_1",
         description: "A test tool",
         parameters: [{ name: "query", type: "string", description: "perform a search" }],
         script: "function invoke(params) { return params; }",
@@ -64,7 +66,8 @@
       }.to change(AiTool, :count).by(1)
 
       expect(response).to have_http_status(:created)
-      expect(response.parsed_body["ai_tool"]["name"]).to eq("Test Tool")
+      expect(response.parsed_body["ai_tool"]["name"]).to eq("Test Tool 1")
+      expect(response.parsed_body["ai_tool"]["tool_name"]).to eq("test_tool_1")
     end
 
     context "when the parameter is a enum" do