FEATURE: Make tool support polymorphic (#798)

Polymorphic RAG means that we will be able to access RAG fragments both from AiPersona and AiCustomTool

In turn this gives us support for richer RAG implementations.

Author: SamSaffron

app/controllers/discourse_ai/admin/ai_personas_controller.rb

@@ -41,7 +41,7 @@ def show
       def create
         ai_persona = AiPersona.new(ai_persona_params.except(:rag_uploads))
         if ai_persona.save
-          RagDocumentFragment.link_persona_and_uploads(ai_persona, attached_upload_ids)
+          RagDocumentFragment.link_target_and_uploads(ai_persona, attached_upload_ids)
 
           render json: {
                    ai_persona: LocalizedAiPersonaSerializer.new(ai_persona, root: false),
@@ -59,7 +59,7 @@ def create_user
 
       def update
         if @ai_persona.update(ai_persona_params.except(:rag_uploads))
-          RagDocumentFragment.update_persona_uploads(@ai_persona, attached_upload_ids)
+          RagDocumentFragment.update_target_uploads(@ai_persona, attached_upload_ids)
 
           render json: LocalizedAiPersonaSerializer.new(@ai_persona, root: false)
         else

app/jobs/regular/digest_rag_upload.rb

@@ -9,17 +9,24 @@ class DigestRagUpload < ::Jobs::Base
     # TODO(roman): Add a way to automatically recover from errors, resulting in unindexed uploads.
     def execute(args)
       return if (upload = Upload.find_by(id: args[:upload_id])).nil?
-      return if (ai_persona = AiPersona.find_by(id: args[:ai_persona_id])).nil?
+
+      target_type = args[:target_type]
+      target_id = args[:target_id]
+
+      return if !target_type || !target_id
+
+      target = target_type.constantize.find_by(id: target_id)
+      return if !target
 
       truncation = DiscourseAi::Embeddings::Strategies::Truncation.new
       vector_rep =
         DiscourseAi::Embeddings::VectorRepresentations::Base.current_representation(truncation)
 
       tokenizer = vector_rep.tokenizer
-      chunk_tokens = ai_persona.rag_chunk_tokens
-      overlap_tokens = ai_persona.rag_chunk_overlap_tokens
+      chunk_tokens = target.rag_chunk_tokens
+      overlap_tokens = target.rag_chunk_overlap_tokens
 
-      fragment_ids = RagDocumentFragment.where(ai_persona: ai_persona, upload: upload).pluck(:id)
+      fragment_ids = RagDocumentFragment.where(target: target, upload: upload).pluck(:id)
 
       # Check if this is the first time we process this upload.
       if fragment_ids.empty?
@@ -39,7 +46,7 @@ def execute(args)
             overlap_tokens: overlap_tokens,
           ) do |chunk, metadata|
             fragment_ids << RagDocumentFragment.create!(
-              ai_persona: ai_persona,
+              target: target,
               fragment: chunk,
               fragment_number: idx + 1,
               upload: upload,

app/jobs/regular/generate_rag_embeddings.rb

@@ -17,10 +17,10 @@ def execute(args)
       fragments.map { |fragment| vector_rep.generate_representation_from(fragment) }
 
       last_fragment = fragments.last
-      ai_persona = last_fragment.ai_persona
+      target = last_fragment.target
       upload = last_fragment.upload
 
-      indexing_status = RagDocumentFragment.indexing_status(ai_persona, [upload])[upload.id]
+      indexing_status = RagDocumentFragment.indexing_status(target, [upload])[upload.id]
       RagDocumentFragment.publish_status(upload, indexing_status)
     end
   end

app/models/ai_persona.rb

@@ -20,19 +20,15 @@ class AiPersona < ActiveRecord::Base
   validates :rag_chunk_tokens, numericality: { greater_than: 0, maximum: 50_000 }
   validates :rag_chunk_overlap_tokens, numericality: { greater_than: -1, maximum: 200 }
   validates :rag_conversation_chunks, numericality: { greater_than: 0, maximum: 1000 }
+  has_many :rag_document_fragments, dependent: :destroy, as: :target
 
   belongs_to :created_by, class_name: "User"
   belongs_to :user
 
   has_many :upload_references, as: :target, dependent: :destroy
   has_many :uploads, through: :upload_references
 
-  has_many :rag_document_fragment, dependent: :destroy
-
-  has_many :rag_document_fragments, through: :ai_persona_rag_document_fragments
-
   before_destroy :ensure_not_system
-
   before_update :regenerate_rag_fragments
 
   def self.persona_cache
@@ -230,7 +226,7 @@ def create_user!
 
   def regenerate_rag_fragments
     if rag_chunk_tokens_changed? || rag_chunk_overlap_tokens_changed?
-      RagDocumentFragment.where(ai_persona: self).delete_all
+      RagDocumentFragment.where(target: self).delete_all
     end
   end
 

app/models/rag_document_fragment.rb

@@ -1,32 +1,40 @@
 # frozen_string_literal: true
 
 class RagDocumentFragment < ActiveRecord::Base
+  # TODO Jan 2025 - remove
+  self.ignored_columns = %i[ai_persona_id]
+
   belongs_to :upload
-  belongs_to :ai_persona
+  belongs_to :target, polymorphic: true
 
   class << self
-    def link_persona_and_uploads(persona, upload_ids)
-      return if persona.blank?
+    def link_target_and_uploads(target, upload_ids)
+      return if target.blank?
       return if upload_ids.blank?
       return if !SiteSetting.ai_embeddings_enabled?
 
-      UploadReference.ensure_exist!(upload_ids: upload_ids, target: persona)
+      UploadReference.ensure_exist!(upload_ids: upload_ids, target: target)
 
       upload_ids.each do |upload_id|
-        Jobs.enqueue(:digest_rag_upload, ai_persona_id: persona.id, upload_id: upload_id)
+        Jobs.enqueue(
+          :digest_rag_upload,
+          target_id: target.id,
+          target_type: target.class.to_s,
+          upload_id: upload_id,
+        )
       end
     end
 
-    def update_persona_uploads(persona, upload_ids)
-      return if persona.blank?
+    def update_target_uploads(target, upload_ids)
+      return if target.blank?
       return if !SiteSetting.ai_embeddings_enabled?
 
       if upload_ids.blank?
-        RagDocumentFragment.where(ai_persona: persona).destroy_all
-        UploadReference.where(target: persona).destroy_all
+        RagDocumentFragment.where(target: target).destroy_all
+        UploadReference.where(target: target).destroy_all
       else
-        RagDocumentFragment.where(ai_persona: persona).where.not(upload_id: upload_ids).destroy_all
-        link_persona_and_uploads(persona, upload_ids)
+        RagDocumentFragment.where(target: target).where.not(upload_id: upload_ids).destroy_all
+        link_target_and_uploads(target, upload_ids)
       end
     end
 
@@ -37,18 +45,25 @@ def indexing_status(persona, uploads)
 
       embeddings_table = vector_rep.rag_fragments_table_name
 
-      results = DB.query(<<~SQL, persona_id: persona.id, upload_ids: uploads.map(&:id))
+      results =
+        DB.query(
+          <<~SQL,
         SELECT
           uploads.id,
           SUM(CASE WHEN (rdf.upload_id IS NOT NULL) THEN 1 ELSE 0 END) AS total,
           SUM(CASE WHEN (eft.rag_document_fragment_id IS NOT NULL) THEN 1 ELSE 0 END) as indexed,
           SUM(CASE WHEN (rdf.upload_id IS NOT NULL AND eft.rag_document_fragment_id IS NULL) THEN 1 ELSE 0 END) as left
         FROM uploads
-        LEFT OUTER JOIN rag_document_fragments rdf ON uploads.id = rdf.upload_id AND rdf.ai_persona_id = :persona_id
+        LEFT OUTER JOIN rag_document_fragments rdf ON uploads.id = rdf.upload_id AND rdf.target_id = :target_id
+          AND rdf.target_type = :target_type
         LEFT OUTER JOIN #{embeddings_table} eft ON rdf.id = eft.rag_document_fragment_id
         WHERE uploads.id IN (:upload_ids)
         GROUP BY uploads.id
       SQL
+          target_id: persona.id,
+          target_type: persona.class.to_s,
+          upload_ids: uploads.map(&:id),
+        )
 
       results.reduce({}) do |acc, r|
         acc[r.id] = { total: r.total, indexed: r.indexed, left: r.left }
@@ -78,4 +93,10 @@ def publish_status(upload, status)
 #  created_at      :datetime         not null
 #  updated_at      :datetime         not null
 #  metadata        :text
+#  target_id       :integer
+#  target_type     :string(800)
+#
+# Indexes
+#
+#  index_rag_document_fragments_on_target_type_and_target_id  (target_type,target_id)
 #

db/migrate/20240912052713_add_target_to_rag_document_fragment.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class AddTargetToRagDocumentFragment < ActiveRecord::Migration[7.1]
+  def change
+    add_column :rag_document_fragments, :target_id, :integer, null: true
+    add_column :rag_document_fragments, :target_type, :string, limit: 800, null: true
+    add_index :rag_document_fragments, %i[target_type target_id]
+  end
+end

db/post_migrate/20240912055831_drop_persona_id_from_rag_document_fragments.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+class DropPersonaIdFromRagDocumentFragments < ActiveRecord::Migration[7.1]
+  def change
+    execute <<~SQL
+      UPDATE rag_document_fragments
+      SET
+        target_type = 'AiPersona',
+        target_id = ai_persona_id
+      WHERE ai_persona_id IS NOT NULL
+    SQL
+
+    # unlikely but lets be safe
+    execute <<~SQL
+      DELETE FROM rag_document_fragments
+      WHERE target_id IS NULL OR target_type IS NULL
+    SQL
+
+    remove_column :rag_document_fragments, :ai_persona_id
+    change_column_null :rag_document_fragments, :target_id, false
+    change_column_null :rag_document_fragments, :target_type, false
+  end
+end

lib/ai_bot/personas/persona.rb

@@ -288,7 +288,8 @@ def rag_fragments_prompt(conversation_context, llm:, user:)
           candidate_fragment_ids =
             vector_rep.asymmetric_rag_fragment_similarity_search(
               interactions_vector,
-              persona_id: id,
+              target_type: "AiPersona",
+              target_id: id,
               limit:
                 (
                   if reranker.reranker_configured?

lib/embeddings/vector_representations/base.rb

@@ -280,7 +280,8 @@ def asymmetric_posts_similarity_search(raw_vector, limit:, offset:, return_dista
 
         def asymmetric_rag_fragment_similarity_search(
           raw_vector,
-          persona_id:,
+          target_id:,
+          target_type:,
           limit:,
           offset:,
           return_distance: false
@@ -299,14 +300,16 @@ def asymmetric_rag_fragment_similarity_search(
             WHERE
               model_id = #{id} AND
               strategy_id = #{@strategy.id} AND
-              rdf.ai_persona_id = :persona_id
+              rdf.target_id = :target_id AND
+              rdf.target_type = :target_type
             ORDER BY
               embeddings::halfvec(#{dimensions}) #{pg_function} '[:query_embedding]'::halfvec(#{dimensions})
             LIMIT :limit
             OFFSET :offset
           SQL
               query_embedding: raw_vector,
-              persona_id: persona_id,
+              target_id: target_id,
+              target_type: target_type,
               limit: limit,
               offset: offset,
             )

spec/jobs/regular/digest_rag_upload_spec.rb

@@ -41,7 +41,11 @@
         # be explicit here about chunking strategy
         persona.update!(rag_chunk_tokens: 100, rag_chunk_overlap_tokens: 10)
 
-        described_class.new.execute(upload_id: upload_with_metadata.id, ai_persona_id: persona.id)
+        described_class.new.execute(
+          upload_id: upload_with_metadata.id,
+          target_id: persona.id,
+          target_type: persona.class.to_s,
+        )
 
         parsed = +""
         first = true
@@ -66,7 +70,11 @@
       before { File.expects(:open).returns(document_file) }
 
       it "splits an upload into chunks" do
-        subject.execute(upload_id: upload.id, ai_persona_id: persona.id)
+        subject.execute(
+          upload_id: upload.id,
+          target_id: persona.id,
+          target_type: persona.class.to_s,
+        )
 
         created_fragment = RagDocumentFragment.last
 
@@ -76,19 +84,23 @@
       end
 
       it "queue jobs to generate embeddings for each fragment" do
-        expect { subject.execute(upload_id: upload.id, ai_persona_id: persona.id) }.to change(
-          Jobs::GenerateRagEmbeddings.jobs,
-          :size,
-        ).by(1)
+        expect {
+          subject.execute(
+            upload_id: upload.id,
+            target_id: persona.id,
+            target_type: persona.class.to_s,
+          )
+        }.to change(Jobs::GenerateRagEmbeddings.jobs, :size).by(1)
       end
     end
 
     it "doesn't generate new fragments if we already processed the upload" do
-      Fabricate(:rag_document_fragment, upload: upload, ai_persona: persona)
-      previous_count = RagDocumentFragment.where(upload: upload, ai_persona: persona).count
+      Fabricate(:rag_document_fragment, upload: upload, target: persona)
 
-      subject.execute(upload_id: upload.id, ai_persona_id: persona.id)
-      updated_count = RagDocumentFragment.where(upload: upload, ai_persona: persona).count
+      previous_count = RagDocumentFragment.where(upload: upload, target: persona).count
+
+      subject.execute(upload_id: upload.id, target_id: persona.id, target_type: persona.class.to_s)
+      updated_count = RagDocumentFragment.where(upload: upload, target: persona).count
 
       expect(updated_count).to eq(previous_count)
     end