Skip to content

Commit

Permalink
FEATURE: Tool name validation (#842)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
nvh0412 authored Feb 7, 2025
1 parent 551f674 commit b60926c
Show file tree
Hide file tree
Showing 18 changed files with 225 additions and 15 deletions.
7 changes: 6 additions & 1 deletion app/controllers/discourse_ai/admin/ai_personas_controller.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down
1 change: 1 addition & 0 deletions app/controllers/discourse_ai/admin/ai_tools_controller.rb
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ def ai_tool_params
.require(:ai_tool)
.permit(
:name,
:tool_name,
:description,
:script,
:summary,
Expand Down
44 changes: 44 additions & 0 deletions app/models/ai_persona.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down
35 changes: 29 additions & 6 deletions app/models/ai_tool.rb
Original file line number Diff line number Diff line change
@@ -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 }
Expand All @@ -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: {})
Expand Down Expand Up @@ -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" },
Expand All @@ -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: [
{
Expand Down Expand Up @@ -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: [
{
Expand Down Expand Up @@ -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: [
Expand Down Expand Up @@ -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
#
1 change: 1 addition & 0 deletions app/serializers/ai_custom_tool_serializer.rb
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
class AiCustomToolSerializer < ApplicationSerializer
attributes :id,
:name,
:tool_name,
:description,
:summary,
:parameters,
Expand Down
1 change: 1 addition & 0 deletions assets/javascripts/discourse/admin/models/ai-tool.js
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import RestModel from "discourse/models/rest";
const CREATE_ATTRIBUTES = [
"id",
"name",
"tool_name",
"description",
"parameters",
"script",
Expand Down
18 changes: 18 additions & 0 deletions assets/javascripts/discourse/components/ai-tool-editor.gjs
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ export default class AiToolEditor extends Component {
try {
const data = this.editingModel.getProperties(
"name",
"tool_name",
"description",
"parameters",
"script",
Expand Down Expand Up @@ -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
Expand Down
6 changes: 4 additions & 2 deletions config/locales/client.en.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
3 changes: 3 additions & 0 deletions config/locales/server.en.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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"
Expand Down
23 changes: 23 additions & 0 deletions db/migrate/20241020010245_add_tool_name_to_ai_tools.rb
Original file line number Diff line number Diff line change
@@ -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
25 changes: 25 additions & 0 deletions db/migrate/20241023041242_add_unique_constraint_to_ai_tools.rb
Original file line number Diff line number Diff line change
@@ -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
1 change: 0 additions & 1 deletion lib/ai_bot/personas/general.rb
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ def tools
Tools::Google,
Tools::Image,
Tools::Read,
Tools::Image,
Tools::ListCategories,
Tools::ListTags,
]
Expand Down
5 changes: 4 additions & 1 deletion lib/ai_bot/tools/custom.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
10 changes: 10 additions & 0 deletions spec/fabricators/ai_tool_fabricator.rb
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions spec/lib/modules/ai_bot/playground_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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" }],
Expand Down
47 changes: 47 additions & 0 deletions spec/models/ai_persona_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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!(
Expand Down
Loading

0 comments on commit b60926c

Please sign in to comment.