FEATURE: RAG search within tools

This implements the back end for RAG search

Author: SamSaffron

app/jobs/regular/digest_rag_upload.rb

@@ -126,7 +126,9 @@ def chunk_document(file:, tokenizer:, chunk_tokens:, overlap_tokens:)
 
         while overlap_token_ids.present?
           begin
-            overlap = tokenizer.decode(overlap_token_ids) + split_char
+            padding = split_char
+            padding = " " if padding.empty?
+            overlap = tokenizer.decode(overlap_token_ids) + padding
             break if overlap.encoding == Encoding::UTF_8
           rescue StandardError
             # it is possible that we truncated mid char
@@ -135,7 +137,7 @@ def chunk_document(file:, tokenizer:, chunk_tokens:, overlap_tokens:)
         end
 
         # remove first word it is probably truncated
-        overlap = overlap.split(" ", 2).last
+        overlap = overlap.split(/\s/, 2).last.to_s.lstrip
       end
     end
 

app/models/ai_tool.rb

@@ -7,6 +7,7 @@ class AiTool < ActiveRecord::Base
   validates :script, presence: true, length: { maximum: 100_000 }
   validates :created_by_id, presence: true
   belongs_to :created_by, class_name: "User"
+  has_many :rag_document_fragments, dependent: :destroy, as: :target
 
   def signature
     { name: name, description: description, parameters: parameters.map(&:symbolize_keys) }
@@ -173,14 +174,16 @@ def self.presets
 #
 # Table name: ai_tools
 #
-#  id            :bigint           not null, primary key
-#  name          :string           not null
-#  description   :string           not null
-#  summary       :string           not null
-#  parameters    :jsonb            not null
-#  script        :text             not null
-#  created_by_id :integer          not null
-#  enabled       :boolean          default(TRUE), not null
-#  created_at    :datetime         not null
-#  updated_at    :datetime         not null
+#  id                       :bigint           not null, primary key
+#  name                     :string           not null
+#  description              :string           not null
+#  summary                  :string           not null
+#  parameters               :jsonb            not null
+#  script                   :text             not null
+#  created_by_id            :integer          not null
+#  enabled                  :boolean          default(TRUE), not null
+#  created_at               :datetime         not null
+#  updated_at               :datetime         not null
+#  rag_chunk_tokens         :integer          default(374), not null
+#  rag_chunk_overlap_tokens :integer          default(10), not null
 #

db/migrate/20240913054440_add_rag_columns_to_ai_tools.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+class AddRagColumnsToAiTools < ActiveRecord::Migration[7.1]
+  def change
+    add_column :ai_tools, :rag_chunk_tokens, :integer, null: false, default: 374
+    add_column :ai_tools, :rag_chunk_overlap_tokens, :integer, null: false, default: 10
+  end
+end

lib/ai_bot/tool_runner.rb

@@ -35,6 +35,7 @@ def mini_racer_context
               )
             attach_truncate(ctx)
             attach_http(ctx)
+            attach_index(ctx)
             ctx.eval(framework_script)
             ctx
           end
@@ -50,6 +51,10 @@ def framework_script
         const llm = {
           truncate: _llm_truncate,
         };
+
+        const index = {
+          search: _index_search,
+        }
         function details() { return ""; };
       JS
       end
@@ -105,13 +110,73 @@ def invoke
 
       private
 
+      MAX_FRAGMENTS = 200
+
+      def rag_search(query, filenames: nil, limit: 10)
+        limit = limit.to_i
+        return [] if limit < 1
+        limit = [MAX_FRAGMENTS, limit].min
+
+        upload_refs = UploadReference.where(target_id: tool.id, target_type: "AiTool").pluck(:upload_id)
+
+        if filenames
+          upload_refs = Upload.where(id: upload_refs).where(original_filename: filenames).pluck(:id)
+        end
+
+        if upload_refs.empty?
+          return []
+        end
+
+        strategy = DiscourseAi::Embeddings::Strategies::Truncation.new
+        vector_rep =
+          DiscourseAi::Embeddings::VectorRepresentations::Base.current_representation(strategy)
+        query_vector = vector_rep.vector_from(query)
+        fragment_ids =
+          vector_rep.asymmetric_rag_fragment_similarity_search(
+            query_vector,
+            target_type: "AiTool",
+            target_id: tool.id,
+            limit: limit,
+            offset: 0
+          )
+        fragments =
+          RagDocumentFragment.where(id: fragment_ids, upload_id: upload_refs).pluck(
+            :id,
+            :fragment,
+            :metadata,
+          )
+
+        mapped = {}
+        fragments.each do |id, fragment, metadata|
+          mapped[id] = { fragment: fragment, metadata: metadata }
+        end
+
+        fragment_ids.take(limit).map { |fragment_id| mapped[fragment_id] }
+      end
+
       def attach_truncate(mini_racer_context)
         mini_racer_context.attach(
           "_llm_truncate",
           ->(text, length) { @llm.tokenizer.truncate(text, length) },
         )
       end
 
+      def attach_index(mini_racer_context)
+        mini_racer_context.attach(
+          "_index_search",
+          ->(query, options) do
+            begin
+              self.running_attached_function = true
+              options ||= {}
+              options = options.symbolize_keys
+              self.rag_search(query, **options)
+            ensure
+              self.running_attached_function = false
+            end
+          end,
+        )
+      end
+
       def attach_http(mini_racer_context)
         mini_racer_context.attach(
           "_http_get",

spec/models/ai_tool_spec.rb

@@ -4,14 +4,21 @@
   fab!(:llm_model) { Fabricate(:llm_model, name: "claude-2") }
   let(:llm) { DiscourseAi::Completions::Llm.proxy("custom:#{llm_model.id}") }
 
-  def create_tool(parameters: nil, script: nil)
+  def create_tool(
+    parameters: nil,
+    script: nil,
+    rag_chunk_tokens: nil,
+    rag_chunk_overlap_tokens: nil
+  )
     AiTool.create!(
       name: "test",
       description: "test",
       parameters: parameters || [{ name: "query", type: "string", desciption: "perform a search" }],
       script: script || "function invoke(params) { return params; }",
       created_by_id: 1,
       summary: "Test tool summary",
+      rag_chunk_tokens: rag_chunk_tokens || 374,
+      rag_chunk_overlap_tokens: rag_chunk_overlap_tokens || 10,
     )
   end
 
@@ -193,4 +200,74 @@ def create_tool(parameters: nil, script: nil)
     result = runner.invoke
     expect(result[:error]).to eq("Script terminated due to timeout")
   end
+
+  context "when defining RAG fragments" do
+    before do
+      SiteSetting.authorized_extensions = "txt"
+      SiteSetting.ai_embeddings_enabled = true
+      SiteSetting.ai_embeddings_discourse_service_api_endpoint = "http://test.com"
+      SiteSetting.ai_embeddings_model = "bge-large-en"
+
+      Jobs.run_immediately!
+    end
+
+    def create_upload(content, filename)
+      upload = nil
+      Tempfile.create(filename) do |file|
+        file.write(content)
+        file.rewind
+
+        upload = UploadCreator.new(file, filename).create_for(Discourse.system_user.id)
+      end
+      upload
+    end
+
+    def stub_embeddings
+      # this is a trick, we get ever increasing embeddings, this gives us in turn
+      # 100% consistent search results
+      @counter = 0
+      stub_request(:post, "http://test.com/api/v1/classify").to_return(
+        status: 200,
+        body: lambda { |req| ([@counter += 1] * 1024).to_json },
+        headers: {
+        },
+      )
+    end
+
+    it "allows search within uploads" do
+      stub_embeddings
+
+      upload1 = create_upload(<<~TXT, "test.txt")
+        1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30
+      TXT
+
+      upload2 = create_upload(<<~TXT, "test.txt")
+        30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50
+      TXT
+
+      tool = create_tool(rag_chunk_tokens: 10, rag_chunk_overlap_tokens: 4, script: <<~JS)
+        function invoke(params) {
+          let result1 = index.search("testing a search", { limit: 1 });
+          let result2 = index.search("testing another search", { limit: 3, filenames: ["test.txt"] });
+
+          return [result1, result2];
+        }
+      JS
+
+      RagDocumentFragment.link_target_and_uploads(tool, [upload1.id, upload2.id])
+
+      result = tool.runner({}, llm: nil, bot_user: nil, context: {}).invoke
+
+      expected = [
+        [{ "fragment" => "44 45 46 47 48 49 50", "metadata" => nil }],
+        [
+          { "fragment" => "44 45 46 47 48 49 50", "metadata" => nil },
+          { "fragment" => "36 37 38 39 40 41 42 43 44 45", "metadata" => nil },
+          { "fragment" => "30 31 32 33 34 35 36 37", "metadata" => nil },
+        ],
+      ]
+
+      expect(result).to eq(expected)
+    end
+  end
 end