FIX: testing tool was not showing rag results (#867)

This changeset contains 4 fixes:

1. We were allowing running tests on unsaved tools,
this is problematic cause uploads are not yet associated or indexed
leading to confusing results. We now only show the test button when
tool is saved.


2. We were not properly scoping rag document fragements, this
meant that personas and ai tools could get results from other
unrelated tools, just to be filtered out later


3. index.search showed options as "optional" but implementation
required the second option

4. When testing tools searching through document fragments was
not working at all cause we did not properly load the tool

Author: SamSaffron

app/controllers/discourse_ai/admin/ai_tools_controller.rb

@@ -5,7 +5,7 @@ module Admin
     class AiToolsController < ::Admin::AdminController
       requires_plugin ::DiscourseAi::PLUGIN_NAME
 
-      before_action :find_ai_tool, only: %i[show update destroy]
+      before_action :find_ai_tool, only: %i[test show update destroy]
 
       def index
         ai_tools = AiTool.all
@@ -17,7 +17,7 @@ def show
       end
 
       def create
-        ai_tool = AiTool.new(ai_tool_params.except(:rag_uploads))
+        ai_tool = AiTool.new(ai_tool_params)
         ai_tool.created_by_id = current_user.id
 
         if ai_tool.save
@@ -29,7 +29,7 @@ def create
       end
 
       def update
-        if @ai_tool.update(ai_tool_params.except(:rag_uploads))
+        if @ai_tool.update(ai_tool_params)
           RagDocumentFragment.update_target_uploads(@ai_tool, attached_upload_ids)
           render_serialized(@ai_tool, AiCustomToolSerializer)
         else
@@ -46,18 +46,13 @@ def destroy
       end
 
       def test
-        if params[:id].present?
-          ai_tool = AiTool.find(params[:id])
-        else
-          ai_tool = AiTool.new(ai_tool_params.except(:rag_uploads))
-        end
-
+        @ai_tool.assign_attributes(ai_tool_params) if params[:ai_tool]
         parameters = params[:parameters].to_unsafe_h
 
         # we need an llm so we have a tokenizer
         # but will do without if none is available
         llm = LlmModel.first&.to_llm
-        runner = ai_tool.runner(parameters, llm: llm, bot_user: current_user, context: {})
+        runner = @ai_tool.runner(parameters, llm: llm, bot_user: current_user, context: {})
         result = runner.invoke
 
         if result.is_a?(Hash) && result[:error]
@@ -74,24 +69,27 @@ def test
       private
 
       def attached_upload_ids
-        ai_tool_params[:rag_uploads].to_a.map { |h| h[:id] }
+        params[:ai_tool][:rag_uploads].to_a.map { |h| h[:id] }
       end
 
       def find_ai_tool
-        @ai_tool = AiTool.find(params[:id])
+        @ai_tool = AiTool.find(params[:id].to_i)
       end
 
       def ai_tool_params
-        params.require(:ai_tool).permit(
-          :name,
-          :description,
-          :script,
-          :summary,
-          :rag_chunk_tokens,
-          :rag_chunk_overlap_tokens,
-          rag_uploads: [:id],
-          parameters: [:name, :type, :description, :required, enum: []],
-        )
+        params
+          .require(:ai_tool)
+          .permit(
+            :name,
+            :description,
+            :script,
+            :summary,
+            :rag_chunk_tokens,
+            :rag_chunk_overlap_tokens,
+            rag_uploads: [:id],
+            parameters: [:name, :type, :description, :required, enum: []],
+          )
+          .except(:rag_uploads)
       end
     end
   end

assets/javascripts/discourse/components/ai-tool-editor.gjs

@@ -232,11 +232,13 @@ export default class AiToolEditor extends Component {
         {{/if}}
 
         <div class="control-group ai-tool-editor__action_panel">
-          <DButton
-            @action={{this.openTestModal}}
-            @label="discourse_ai.tools.test"
-            class="ai-tool-editor__test-button"
-          />
+          {{#unless @model.isNew}}
+            <DButton
+              @action={{this.openTestModal}}
+              @label="discourse_ai.tools.test"
+              class="ai-tool-editor__test-button"
+            />
+          {{/unless}}
 
           <DButton
             @action={{this.save}}

assets/javascripts/discourse/components/modal/ai-tool-test-modal.gjs

@@ -25,7 +25,7 @@ export default class AiToolTestModal extends Component {
     this.isLoading = true;
     try {
       const response = await ajax(
-        "/admin/plugins/discourse-ai/ai-tools/test.json",
+        `/admin/plugins/discourse-ai/ai-tools/${this.args.model.tool.id}/test.json`,
         {
           type: "POST",
           data: JSON.stringify({

config/routes.rb

@@ -55,7 +55,9 @@
       only: %i[index create show update destroy],
       path: "ai-tools",
       controller: "discourse_ai/admin/ai_tools",
-    ) { post :test, on: :collection }
+    )
+
+    post "/ai-tools/:id/test", to: "discourse_ai/admin/ai_tools#test"
 
     post "/ai-personas/:id/create-user", to: "discourse_ai/admin/ai_personas#create_user"
 

lib/ai_bot/tool_runner.rb

@@ -149,6 +149,7 @@ def rag_search(query, filenames: nil, limit: 10)
             limit: limit,
             offset: 0,
           )
+
         fragments =
           RagDocumentFragment.where(id: fragment_ids, upload_id: upload_refs).pluck(
             :id,
@@ -174,8 +175,9 @@ def attach_truncate(mini_racer_context)
       def attach_index(mini_racer_context)
         mini_racer_context.attach(
           "_index_search",
-          ->(query, options) do
+          ->(*params) do
             begin
+              query, options = params
               self.running_attached_function = true
               options ||= {}
               options = options.symbolize_keys

lib/embeddings/vector_representations/base.rb

@@ -205,7 +205,10 @@ def asymmetric_rag_fragment_similarity_search(
                   FROM
                     #{rag_fragments_table_name}
                   INNER JOIN
-                    rag_document_fragments ON rag_document_fragments.id = rag_document_fragment_id
+                    rag_document_fragments ON
+                      rag_document_fragments.id = rag_document_fragment_id AND
+                      rag_document_fragments.target_id = :target_id AND
+                      rag_document_fragments.target_type = :target_type
                   WHERE
                     model_id = #{id} AND strategy_id = #{@strategy.id}
                   ORDER BY

spec/requests/admin/ai_tools_controller_spec.rb

@@ -137,9 +137,8 @@
 
   describe "#test" do
     it "runs an existing tool and returns the result" do
-      post "/admin/plugins/discourse-ai/ai-tools/test.json",
+      post "/admin/plugins/discourse-ai/ai-tools/#{ai_tool.id}/test.json",
            params: {
-             id: ai_tool.id,
              parameters: {
                input: "Hello, World!",
              },
@@ -149,43 +148,36 @@
       expect(response.parsed_body["output"]).to eq("input" => "Hello, World!")
     end
 
-    it "runs a new unsaved tool and returns the result" do
-      post "/admin/plugins/discourse-ai/ai-tools/test.json",
+    it "accept changes to the ai_tool parameters that redefine stuff" do
+      post "/admin/plugins/discourse-ai/ai-tools/#{ai_tool.id}/test.json",
            params: {
              ai_tool: {
-               name: "New Tool",
-               description: "A new test tool",
-               script: "function invoke(params) { return 'New test result: ' + params.input; }",
-               parameters: [
-                 { name: "input", type: "string", description: "Input for the new test tool" },
-               ],
+               script: "function invoke(params) { return 'hi there'; }",
              },
              parameters: {
-               input: "Test input",
+               input: "Hello, World!",
              },
            }
 
       expect(response.status).to eq(200)
-      expect(response.parsed_body["output"]).to eq("New test result: Test input")
+      expect(response.parsed_body["output"]).to eq("hi there")
     end
 
     it "returns an error for invalid tool_id" do
-      post "/admin/plugins/discourse-ai/ai-tools/test.json",
+      post "/admin/plugins/discourse-ai/ai-tools/-1/test.json",
            params: {
-             id: -1,
              parameters: {
                input: "Hello, World!",
              },
            }
 
-      expect(response.status).to eq(400)
-      expect(response.parsed_body["errors"]).to include("Couldn't find AiTool with 'id'=-1")
+      expect(response.status).to eq(404)
     end
 
     it "handles exceptions during tool execution" do
       ai_tool.update!(script: "function invoke(params) { throw new Error('Test error'); }")
 
-      post "/admin/plugins/discourse-ai/ai-tools/test.json",
+      post "/admin/plugins/discourse-ai/ai-tools/#{ai_tool.id}/test.json",
            params: {
              id: ai_tool.id,
              parameters: {

spec/system/ai_bot/tool_spec.rb

@@ -49,7 +49,8 @@ def ensure_can_run_test
     expect(page.first(".parameter-row__required-toggle").checked?).to eq(true)
     expect(page.first(".parameter-row__enum-toggle").checked?).to eq(false)
 
-    ensure_can_run_test
+    # not allowed to test yet
+    expect(page).not_to have_button(".ai-tool-editor__test-button")
 
     expect(page).not_to have_button(".ai-tool-editor__delete")
     find(".ai-tool-editor__save").click