From f328b81c78e7ab95cda3cb9d33ccc85beb26c16f Mon Sep 17 00:00:00 2001
From: Roman Rizzi <roman@discourse.org>
Date: Tue, 16 Jul 2024 14:23:17 -0300
Subject: [PATCH] FIX: Make sure custom tool enums follow json-schema. (#718)

Enums didn't work as expected because we the dialect couldn't translate
them correctly. It doesn't understand what "enum_values" is.
---
 .../discourse_ai/admin/ai_tools_controller.rb |  2 +-
 .../discourse/admin/models/ai-tool.js         |  6 ++-
 .../discourse/components/ai-tool-editor.gjs   |  9 +---
 .../components/ai-tool-parameter-editor.gjs   | 25 ++++-----
 config/locales/client.en.yml                  |  1 +
 lib/completions/dialects/open_ai_tools.rb     |  5 +-
 .../admin/ai_tools_controller_spec.rb         | 52 ++++++++++++++++++-
 7 files changed, 76 insertions(+), 24 deletions(-)

diff --git a/app/controllers/discourse_ai/admin/ai_tools_controller.rb b/app/controllers/discourse_ai/admin/ai_tools_controller.rb
index a0341089..4eb53578 100644
--- a/app/controllers/discourse_ai/admin/ai_tools_controller.rb
+++ b/app/controllers/discourse_ai/admin/ai_tools_controller.rb
@@ -81,7 +81,7 @@ module DiscourseAi
           :description,
           :script,
           :summary,
-          parameters: [:name, :type, :description, :required, :enum, enum_values: []],
+          parameters: [:name, :type, :description, :required, enum: []],
         )
       end
     end
diff --git a/assets/javascripts/discourse/admin/models/ai-tool.js b/assets/javascripts/discourse/admin/models/ai-tool.js
index 7f1a0877..5002806e 100644
--- a/assets/javascripts/discourse/admin/models/ai-tool.js
+++ b/assets/javascripts/discourse/admin/models/ai-tool.js
@@ -27,9 +27,13 @@ export default class AiTool extends RestModel {
       attrs.parameters?.map((p) => {
         const parameter = new TrackedObject(p);
 
+        //Backwards-compatibility code.
+        // TODO(roman): Remove aug 2024. Leave only else clause.
         if (parameter.enum_values) {
-          parameter.enumValues = new TrackedArray(parameter.enum_values);
+          parameter.enum = new TrackedArray(parameter.enum_values);
           delete parameter.enum_values;
+        } else {
+          parameter.enum = new TrackedArray(parameter.enum);
         }
 
         return parameter;
diff --git a/assets/javascripts/discourse/components/ai-tool-editor.gjs b/assets/javascripts/discourse/components/ai-tool-editor.gjs
index 73893689..e9c3f70c 100644
--- a/assets/javascripts/discourse/components/ai-tool-editor.gjs
+++ b/assets/javascripts/discourse/components/ai-tool-editor.gjs
@@ -71,13 +71,6 @@ export default class AiToolEditor extends Component {
         "summary"
       );
 
-      for (const p of data.parameters) {
-        if (p.enumValues) {
-          p.enum_values = p.enumValues;
-          delete p.enumValues;
-        }
-      }
-
       await this.args.model.save(data);
 
       this.toasts.success({
@@ -124,7 +117,7 @@ export default class AiToolEditor extends Component {
   <template>
     <BackButton
       @route="adminPlugins.show.discourse-ai-tools"
-      @label="discourse_ai.ai_tool.back"
+      @label="discourse_ai.tools.back"
     />
 
     <form
diff --git a/assets/javascripts/discourse/components/ai-tool-parameter-editor.gjs b/assets/javascripts/discourse/components/ai-tool-parameter-editor.gjs
index 8ee3821e..dda4bde0 100644
--- a/assets/javascripts/discourse/components/ai-tool-parameter-editor.gjs
+++ b/assets/javascripts/discourse/components/ai-tool-parameter-editor.gjs
@@ -7,6 +7,7 @@ import DButton from "discourse/components/d-button";
 import withEventValue from "discourse/helpers/with-event-value";
 import I18n from "discourse-i18n";
 import ComboBox from "select-kit/components/combo-box";
+import and from "truth-helpers/helpers/and";
 
 const PARAMETER_TYPES = [
   { name: "string", id: "string" },
@@ -24,8 +25,7 @@ export default class AiToolParameterEditor extends Component {
         description: "",
         type: "string",
         required: false,
-        enum: false,
-        enumValues: null,
+        enum: null,
       })
     );
   }
@@ -43,26 +43,27 @@ export default class AiToolParameterEditor extends Component {
 
   @action
   toggleEnum(parameter) {
-    parameter.enum = !parameter.enum;
-    if (!parameter.enum) {
-      parameter.enumValues = null;
+    if (parameter.enum) {
+      parameter.enum = null;
+    } else {
+      this.addEnumValue(parameter);
     }
   }
 
   @action
   addEnumValue(parameter) {
-    parameter.enumValues ||= new TrackedArray();
-    parameter.enumValues.push("");
+    parameter.enum ||= new TrackedArray();
+    parameter.enum.push("");
   }
 
   @action
   removeEnumValue(parameter, index) {
-    parameter.enumValues.splice(index, 1);
+    parameter.enum.splice(index, 1);
   }
 
   @action
   updateEnumValue(parameter, index, event) {
-    parameter.enumValues[index] = event.target.value;
+    parameter.enum[index] = event.target.value;
   }
 
   <template>
@@ -100,7 +101,7 @@ export default class AiToolParameterEditor extends Component {
           <label>
             <input
               {{on "input" (fn this.toggleEnum parameter)}}
-              checked={{parameter.enum}}
+              checked={{and parameter.enum parameter.enum.length}}
               type="checkbox"
             />
             {{I18n.t "discourse_ai.tools.parameter_enum"}}
@@ -113,9 +114,9 @@ export default class AiToolParameterEditor extends Component {
           />
         </div>
 
-        {{#if parameter.enum}}
+        {{#if (and parameter.enum parameter.enum.length)}}
           <div class="parameter-enum-values">
-            {{#each parameter.enumValues as |enumValue enumIndex|}}
+            {{#each parameter.enum as |enumValue enumIndex|}}
               <div class="enum-value-row">
                 <input
                   {{on "change" (fn this.updateEnumValue parameter enumIndex)}}
diff --git a/config/locales/client.en.yml b/config/locales/client.en.yml
index 0de84895..039ca3eb 100644
--- a/config/locales/client.en.yml
+++ b/config/locales/client.en.yml
@@ -184,6 +184,7 @@ en:
           remove: "Remove upload"
 
       tools:
+        back: "Back"
         short_title: "Tools"
         new: "New Tool"
         name: "Name"
diff --git a/lib/completions/dialects/open_ai_tools.rb b/lib/completions/dialects/open_ai_tools.rb
index b990e379..73b01146 100644
--- a/lib/completions/dialects/open_ai_tools.rb
+++ b/lib/completions/dialects/open_ai_tools.rb
@@ -18,7 +18,10 @@ module DiscourseAi
                 name = p[:name]
                 memo[:required] << name if p[:required]
 
-                memo[:properties][name] = p.except(:name, :required, :item_type)
+                except = %i[name required item_type]
+                except << :enum if p[:enum].blank?
+
+                memo[:properties][name] = p.except(*except)
 
                 memo[:properties][name][:items] = { type: p[:item_type] } if p[:item_type]
                 memo
diff --git a/spec/requests/admin/ai_tools_controller_spec.rb b/spec/requests/admin/ai_tools_controller_spec.rb
index 42a61e7f..e05b7344 100644
--- a/spec/requests/admin/ai_tools_controller_spec.rb
+++ b/spec/requests/admin/ai_tools_controller_spec.rb
@@ -7,7 +7,15 @@ RSpec.describe DiscourseAi::Admin::AiToolsController do
       name: "Test Tool",
       description: "A test tool",
       script: "function invoke(params) { return params; }",
-      parameters: [{ name: "query", type: "string", description: "perform a search" }],
+      parameters: [
+        {
+          name: "unit",
+          type: "string",
+          description: "the unit of measurement celcius c or fahrenheit f",
+          enum: %w[c f],
+          required: true,
+        },
+      ],
       summary: "Test tool summary",
       created_by_id: -1,
     )
@@ -58,6 +66,27 @@ RSpec.describe DiscourseAi::Admin::AiToolsController do
       expect(response).to have_http_status(:created)
       expect(response.parsed_body["ai_tool"]["name"]).to eq("Test Tool")
     end
+
+    context "when the parameter is a enum" do
+      it "creates the tool with the correct parameters" do
+        attrs = valid_attributes
+        attrs[:parameters] = [attrs[:parameters].first.merge(enum: %w[c f])]
+
+        expect {
+          post "/admin/plugins/discourse-ai/ai-tools.json",
+               params: { ai_tool: valid_attributes }.to_json,
+               headers: {
+                 "CONTENT_TYPE" => "application/json",
+               }
+        }.to change(AiTool, :count).by(1)
+
+        expect(response).to have_http_status(:created)
+        expect(response.parsed_body.dig("ai_tool", "parameters", 0, "enum")).to contain_exactly(
+          "c",
+          "f",
+        )
+      end
+    end
   end
 
   describe "PUT #update" do
@@ -72,6 +101,27 @@ RSpec.describe DiscourseAi::Admin::AiToolsController do
       expect(response).to be_successful
       expect(ai_tool.reload.name).to eq("Updated Tool")
     end
+
+    context "when updating an enum parameters" do
+      it "updates the enum fixed values" do
+        put "/admin/plugins/discourse-ai/ai-tools/#{ai_tool.id}.json",
+            params: {
+              ai_tool: {
+                parameters: [
+                  {
+                    name: "unit",
+                    type: "string",
+                    description: "the unit of measurement celcius c or fahrenheit f",
+                    enum: %w[g d],
+                  },
+                ],
+              },
+            }
+
+        expect(response).to be_successful
+        expect(ai_tool.reload.parameters.dig(0, "enum")).to contain_exactly("g", "d")
+      end
+    end
   end
 
   describe "DELETE #destroy" do