Drop usage of #asymmetric_rag_fragment_similarity_search. Don't rely on stubs

romanrizzi · romanrizzi · commit 75ad96757a7a · 2024-12-12T11:51:33.000-03:00
diff --git a/app/models/rag_document_fragment.rb b/app/models/rag_document_fragment.rb
@@ -39,9 +39,7 @@ def update_target_uploads(target, upload_ids)
     end
 
     def indexing_status(persona, uploads)
-      vector_rep = DiscourseAi::Embeddings::VectorRepresentations::Base.current_representation
-
-      embeddings_table = vector_rep.rag_fragments_table_name
+      embeddings_table = DiscourseAi::Embeddings::Schema.for(self).table
 
       results =
         DB.query(
diff --git a/lib/ai_bot/personas/persona.rb b/lib/ai_bot/personas/persona.rb
@@ -320,22 +320,28 @@ def rag_fragments_prompt(conversation_context, llm:, user:)
           interactions_vector = vector_rep.vector_from(consolidated_question)
 
           rag_conversation_chunks = self.class.rag_conversation_chunks
+          search_limit =
+            if reranker.reranker_configured?
+              rag_conversation_chunks * 5
+            else
+              rag_conversation_chunks
+            end
+
+          schema = DiscourseAi::Embeddings::Schema.for(RagDocumentFragment, vector: vector_rep)
 
           candidate_fragment_ids =
-            vector_rep.asymmetric_rag_fragment_similarity_search(
-              interactions_vector,
-              target_type: "AiPersona",
-              target_id: id,
-              limit:
-                (
-                  if reranker.reranker_configured?
-                    rag_conversation_chunks * 5
-                  else
-                    rag_conversation_chunks
-                  end
-                ),
-              offset: 0,
-            )
+            schema
+              .asymmetric_similarity_search(
+                interactions_vector,
+                limit: search_limit,
+                offset: 0,
+              ) { |builder| builder.join(<<~SQL, target_id: id, target_type: "AiPersona") }
+                  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
+                SQL
+              .map(&:rag_document_fragment_id)
 
           fragments =
             RagDocumentFragment.where(upload_id: upload_refs, id: candidate_fragment_ids).pluck(
diff --git a/lib/ai_bot/tool_runner.rb b/lib/ai_bot/tool_runner.rb
@@ -144,13 +144,17 @@ def rag_search(query, filenames: nil, limit: 10)
         vector_rep = DiscourseAi::Embeddings::VectorRepresentations::Base.current_representation
         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,
-          )
+          DiscourseAi::Embeddings::Schema
+            .for(RagDocumentFragment, vector: vector_rep)
+            .asymmetric_similarity_search(query_vector, limit: limit, offset: 0) do |builder|
+              builder.join(<<~SQL, target_id: tool.id, target_type: "AiTool")
+                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
+              SQL
+            end
+            .map(&:rag_document_fragment_id)
 
         fragments =
           RagDocumentFragment.where(id: fragment_ids, upload_id: upload_refs).pluck(
diff --git a/lib/embeddings/schema.rb b/lib/embeddings/schema.rb
@@ -72,7 +72,7 @@ def asymmetric_similarity_search(embedding, limit:, offset:)
             /*where*/
             ORDER BY
               binary_quantize(embeddings)::bit(#{dimensions}) <~> binary_quantize('[:query_embedding]'::halfvec(#{dimensions}))
-            LIMIT :limit * 2
+            LIMIT :candidates_limit
           )
           SELECT
             #{target_column},
@@ -93,7 +93,19 @@ def asymmetric_similarity_search(embedding, limit:, offset:)
 
         yield(builder) if block_given?
 
-        builder.query(query_embedding: embedding, limit: limit, offset: offset)
+        if table == RAG_DOCS_TABLE
+          # A too low limit exacerbates the the recall loss of binary quantization
+          candidates_limit = [limit * 2, 100].max
+        else
+          candidates_limit = limit * 2
+        end
+
+        builder.query(
+          query_embedding: embedding,
+          candidates_limit: candidates_limit,
+          limit: limit,
+          offset: offset,
+        )
       rescue PG::Error => e
         Rails.logger.error("Error #{e} querying embeddings for model #{name}")
         raise MissingEmbeddingError
diff --git a/lib/embeddings/vector_representations/base.rb b/lib/embeddings/vector_representations/base.rb
@@ -132,64 +132,6 @@ def post_id_from_representation(raw_vector)
           SQL
         end
 
-        def asymmetric_rag_fragment_similarity_search(
-          raw_vector,
-          target_id:,
-          target_type:,
-          limit:,
-          offset:,
-          return_distance: false
-        )
-          # A too low limit exacerbates the the recall loss of binary quantization
-          binary_search_limit = [limit * 2, 100].max
-          results =
-            DB.query(
-              <<~SQL,
-                WITH candidates AS (
-                  SELECT
-                    rag_document_fragment_id,
-                    embeddings::halfvec(#{dimensions}) AS embeddings
-                  FROM
-                    #{rag_fragments_table_name}
-                  INNER JOIN
-                    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
-                    binary_quantize(embeddings)::bit(#{dimensions}) <~> binary_quantize('[:query_embedding]'::halfvec(#{dimensions}))
-                  LIMIT :binary_search_limit
-                )
-                SELECT
-                  rag_document_fragment_id,
-                  embeddings::halfvec(#{dimensions}) #{pg_function} '[:query_embedding]'::halfvec(#{dimensions}) AS distance
-                FROM
-                  candidates
-                ORDER BY
-                  embeddings::halfvec(#{dimensions}) #{pg_function} '[:query_embedding]'::halfvec(#{dimensions})
-                LIMIT :limit
-                OFFSET :offset
-              SQL
-              query_embedding: raw_vector,
-              target_id: target_id,
-              target_type: target_type,
-              limit: limit,
-              offset: offset,
-              binary_search_limit: binary_search_limit,
-            )
-
-          if return_distance
-            results.map { |r| [r.rag_document_fragment_id, r.distance] }
-          else
-            results.map(&:rag_document_fragment_id)
-          end
-        rescue PG::Error => e
-          Rails.logger.error("Error #{e} querying embeddings for model #{name}")
-          raise MissingEmbeddingError
-        end
-
         def symmetric_topics_similarity_search(topic)
           DB.query(<<~SQL, topic_id: topic.id).map(&:topic_id)
             WITH le_target AS (
diff --git a/spec/lib/modules/ai_bot/personas/persona_spec.rb b/spec/lib/modules/ai_bot/personas/persona_spec.rb
@@ -373,41 +373,44 @@ def system_prompt
     end
 
     context "when a persona has RAG uploads" do
-      def stub_fragments(limit, expected_limit: nil)
-        candidate_ids = []
-
-        limit.times do |i|
-          candidate_ids << Fabricate(
-            :rag_document_fragment,
-            fragment: "fragment-n#{i}",
-            target_id: ai_persona.id,
-            target_type: "AiPersona",
-            upload: upload,
-          ).id
-        end
+      let(:vector_rep) do
+        DiscourseAi::Embeddings::VectorRepresentations::Base.current_representation
+      end
+      let(:embedding_value) { 0.04381 }
+      let(:prompt_cc_embeddings) { [embedding_value] * vector_rep.dimensions }
+
+      def stub_fragments(fragment_count, persona: ai_persona)
+        schema = DiscourseAi::Embeddings::Schema.for(RagDocumentFragment, vector: vector_rep)
+
+        fragment_count.times do |i|
+          fragment =
+            Fabricate(
+              :rag_document_fragment,
+              fragment: "fragment-n#{i}",
+              target_id: persona.id,
+              target_type: "AiPersona",
+              upload: upload,
+            )
 
-        DiscourseAi::Embeddings::VectorRepresentations::BgeLargeEn
-          .any_instance
-          .expects(:asymmetric_rag_fragment_similarity_search)
-          .with { |args, kwargs| kwargs[:limit] == (expected_limit || limit) }
-          .returns(candidate_ids)
+          # Similarity is determined left-to-right.
+          embeddings = [embedding_value + "0.000#{i}".to_f] * vector_rep.dimensions
+
+          schema.store(fragment, embeddings, "test")
+        end
       end
 
       before do
         stored_ai_persona = AiPersona.find(ai_persona.id)
         UploadReference.ensure_exist!(target: stored_ai_persona, upload_ids: [upload.id])
 
-        context_embedding = [0.049382, 0.9999]
         EmbeddingsGenerationStubs.discourse_service(
           SiteSetting.ai_embeddings_model,
           with_cc.dig(:conversation_context, 0, :content),
-          context_embedding,
+          prompt_cc_embeddings,
         )
       end
 
       context "when persona allows for less fragments" do
-        before { stub_fragments(3) }
-
         it "will only pick 3 fragments" do
           custom_ai_persona =
             Fabricate(
@@ -417,6 +420,8 @@ def stub_fragments(limit, expected_limit: nil)
               allowed_group_ids: [Group::AUTO_GROUPS[:trust_level_0]],
             )
 
+          stub_fragments(3, persona: custom_ai_persona)
+
           UploadReference.ensure_exist!(target: custom_ai_persona, upload_ids: [upload.id])
 
           custom_persona =
@@ -436,13 +441,10 @@ def stub_fragments(limit, expected_limit: nil)
       context "when the reranker is available" do
         before do
           SiteSetting.ai_hugging_face_tei_reranker_endpoint = "https://test.reranker.com"
-
-          # hard coded internal implementation, reranker takes x5 number of chunks
-          stub_fragments(15, expected_limit: 50) # Mimic limit being more than 10 results
+          stub_fragments(15)
         end
 
         it "uses the re-ranker to reorder the fragments and pick the top 10 candidates" do
-          skip "This test is flaky needs to be investigated ordering does not come back as expected"
           expected_reranked = (0..14).to_a.reverse.map { |idx| { index: idx } }
 
           WebMock.stub_request(:post, "https://test.reranker.com/rerank").to_return(
