FEATURE: Use embeddings to find similar topics to the one being composed (#34448)

Co-authored-by: Sam Saffron <sam.saffron@gmail.com>

app/models/topic.rb

@@ -709,6 +709,7 @@ class Topic < ActiveRecord::Base
   MAX_SIMILAR_BODY_LENGTH = 200
   SIMILAR_TOPIC_SEARCH_LIMIT = 10
   SIMILAR_TOPIC_LIMIT = 3
+  SIMILAR_TOPIC_MAX_BLURB_LENGTH = 500
 
   def self.similar_to(title, raw, user = nil)
     return [] if title.blank?
@@ -748,18 +749,66 @@ class Topic < ActiveRecord::Base
       #{excluded_category_ids_sql}
       UNION
       #{CategoryUser.muted_category_ids_query(user, include_direct: true).select("categories.id").to_sql}
-      SQL
+    SQL
 
     candidates =
       Topic
         .visible
         .listable_topics
         .secured(guardian)
-        .joins("JOIN topic_search_data s ON topics.id = s.topic_id")
         .joins("LEFT JOIN categories c ON topics.id = c.topic_id")
-        .where("search_data @@ #{tsquery}")
         .where("c.topic_id IS NULL")
         .where("topics.category_id NOT IN (#{excluded_category_ids_sql})")
+
+    plugin_candidate_ids = []
+    plugin_candidate_ids =
+      DiscoursePluginRegistry.apply_modifier(
+        :similar_topic_candidate_ids,
+        plugin_candidate_ids,
+        title:,
+        raw:,
+        guardian:,
+        candidates:,
+      )
+
+    blurb_sql = "LEFT(p.cooked, #{SIMILAR_TOPIC_MAX_BLURB_LENGTH.to_i}) AS blurb"
+    select_fragment = ->(similarity_expr) do
+      DB.sql_fragment(
+        "topics.*, #{similarity_expr} AS similarity, #{blurb_sql}",
+        title: title,
+        raw: raw,
+      )
+    end
+
+    if plugin_candidate_ids.present? && plugin_candidate_ids.length > 0
+      ids = plugin_candidate_ids.map(&:to_i)
+      candidate_ids =
+        candidates
+          .where(id: ids)
+          .order("array_position(ARRAY[#{ids.join(",")}]::int[], topics.id)")
+          .limit(SIMILAR_TOPIC_LIMIT)
+          .pluck(:id)
+
+      if candidate_ids.present? && candidate_ids.length > 0
+        rank_array_sql = "ARRAY[#{candidate_ids.map(&:to_i).join(",")}]"
+        return(
+          Topic
+            .joins("JOIN posts AS p ON p.topic_id = topics.id AND p.post_number = 1")
+            .where(id: candidate_ids)
+            .select(
+              select_fragment.call(
+                "(array_length(#{rank_array_sql}, 1) - array_position(#{rank_array_sql}, topics.id) + 1)",
+              ),
+            )
+            .order("similarity DESC")
+        )
+      end
+    end
+
+    candidates =
+      candidates
+        .joins("JOIN topic_search_data s ON topics.id = s.topic_id")
+        .where("search_data @@ #{tsquery}")
         .order("ts_rank(search_data, #{tsquery}) DESC")
         .limit(SIMILAR_TOPIC_SEARCH_LIMIT)
 
@@ -776,23 +825,17 @@ class Topic < ActiveRecord::Base
 
     if raw.present?
       similars.select(
-        DB.sql_fragment(
-          "topics.*, similarity(topics.title, :title) + similarity(p.raw, :raw) AS similarity, p.cooked AS blurb",
-          title: title,
-          raw: raw,
-        ),
+        select_fragment.call("similarity(topics.title, :title) + similarity(p.raw, :raw)"),
       ).where(
         "similarity(topics.title, :title) + similarity(p.raw, :raw) > 0.2",
         title: title,
         raw: raw,
       )
     else
-      similars.select(
-        DB.sql_fragment(
-          "topics.*, similarity(topics.title, :title) AS similarity, p.cooked AS blurb",
-          title: title,
-        ),
-      ).where("similarity(topics.title, :title) > 0.2", title: title)
+      similars.select(select_fragment.call("similarity(topics.title, :title)")).where(
+        "similarity(topics.title, :title) > 0.2",
+        title: title,
+      )
     end
   end
 

plugins/discourse-ai/lib/embeddings/entry_point.rb

@@ -72,6 +72,17 @@ module DiscourseAi
           :discourse_ai,
           { search: { actions: %w[discourse_ai/embeddings/embeddings#search] } },
         )
+
+        plugin.register_modifier(:similar_topic_candidate_ids) do |plugin_candidate_ids, args|
+          if DiscourseAi::Embeddings.enabled?
+            query = [args[:title], args[:raw]].join("\n\n")
+
+            DiscourseAi::Embeddings::SemanticSearch
+              .new(args[:guardian])
+              .similar_topic_ids_to(query, candidates: args[:candidates])
+              .each { |similar_topic_id| plugin_candidate_ids << similar_topic_id }
+          end
+        end
       end
     end
   end

plugins/discourse-ai/lib/embeddings/semantic_search.rb

@@ -106,6 +106,25 @@ module DiscourseAi
         guardian.filter_allowed_categories(query_filter_results)
       end
 
+      def similar_topic_ids_to(query, candidates:)
+        return [] if candidates.blank?
+
+        over_selection_limit = ::Topic::SIMILAR_TOPIC_LIMIT * OVER_SELECTION_FACTOR
+        asymmetric = true
+        search_embedding = vector.vector_from(query, asymmetric)
+
+        schema = DiscourseAi::Embeddings::Schema.for(Topic)
+
+        candidate_topic_ids =
+          schema.asymmetric_similarity_search(
+            search_embedding,
+            limit: over_selection_limit,
+            offset: 0,
+          ).map(&:topic_id)
+
+        candidates.where(id: candidate_topic_ids).pluck(:id)
+      end
+
       def quick_search(query)
         max_semantic_results_per_page = 100
         search = Search.new(query, { guardian: guardian })

plugins/discourse-ai/spec/lib/modules/embeddings/entry_point_spec.rb

@@ -39,4 +39,63 @@ describe DiscourseAi::Embeddings::EntryPoint do
       end
     end
   end
+
+  describe "similar_topic_candidate_ids modifier" do
+    # The Distance gap to target increases for each element of topics.
+    def seed_embeddings(topics)
+      schema = DiscourseAi::Embeddings::Schema.for(Topic)
+      base_value = 1
+
+      topics.each_with_index do |t, idx|
+        base_value -= 0.01
+        schema.store(t, [base_value] * embedding_definition.dimensions, "digest")
+      end
+    end
+
+    def stub_query_embedding(query)
+      embedding = [1] * embedding_definition.dimensions
+
+      EmbeddingsGenerationStubs.hugging_face_service(query, embedding)
+    end
+
+    fab!(:category)
+    fab!(:normal_topic_1) { Fabricate(:topic, category: category) }
+    fab!(:normal_topic_2) { Fabricate(:topic, category: category) }
+    fab!(:private_topic) { Fabricate(:private_message_topic) }
+
+    let(:query) { "title\n\nraw" }
+
+    fab!(:embedding_definition)
+
+    before do
+      [normal_topic_1, normal_topic_2, private_topic].each_with_index do |t, idx|
+        Fabricate(
+          :post,
+          topic: t,
+          user: t.user,
+          post_number: 1,
+          raw: "This is a post with raw ##{idx + 1}",
+        )
+      end
+
+      seed_embeddings([normal_topic_1, private_topic])
+      stub_query_embedding(query)
+      SiteSetting.ai_embeddings_enabled = true
+      SiteSetting.ai_embeddings_selected_model = embedding_definition.id
+    end
+
+    it "appends topic IDs" do
+      similar_topics = Topic.similar_to("title", "raw")
+
+      expect(similar_topics.map(&:id)).to contain_exactly(normal_topic_1.id)
+    end
+
+    it "does nothing if embeddings is not enabled" do
+      SiteSetting.ai_embeddings_enabled = false
+
+      similar_topics = Topic.similar_to("title", "raw")
+
+      expect(similar_topics.map(&:id)).to be_empty
+    end
+  end
 end

spec/models/topic_spec.rb

@@ -683,6 +683,64 @@ describe Topic do
       expect(Topic.similar_to("'bad quotes'", "'bad quotes'")).to eq([])
     end
 
+    context "with plugin similar_topic_candidate_ids modifier" do
+      it "uses plugin-provided candidate ids preserving order and respecting limit" do
+        t1 = Fabricate(:topic)
+        t2 = Fabricate(:topic)
+        t3 = Fabricate(:topic)
+        t4 = Fabricate(:topic)
+
+        raws = { t1.id => "raw one", t2.id => "raw two", t3.id => "raw three", t4.id => "raw four" }
+
+        [t1, t2, t3, t4].each do |t|
+          Fabricate(:post, topic: t, user: t.user, post_number: 1, raw: raws[t.id])
+        end
+
+        desired_order = [t3.id, t1.id, t2.id, t4.id]
+
+        plugin_instance = Plugin::Instance.new
+        begin
+          blk =
+            lambda do |candidates, args|
+              expect(args[:title]).to eq("any title")
+              expect(args[:raw]).to eq("any raw")
+              desired_order
+            end
+
+          DiscoursePluginRegistry.register_modifier(
+            plugin_instance,
+            :similar_topic_candidate_ids,
+            &blk
+          )
+
+          results = Topic.similar_to("any title", "any raw")
+
+          # keeping this 3 but test will break if MAX_SIMILAR_TOPICS is changed (by design)
+          expected_ids = desired_order.first(3)
+          expect(results.map(&:id)).to eq(expected_ids)
+
+          # ensure extra selected columns are present and correct
+          results.each_with_index do |topic, idx|
+            # topics.* still present
+            expect(topic).to be_a(Topic)
+            expect(topic.id).to eq(expected_ids[idx])
+
+            # similarity is computed as 3,2,1 for our limited set
+            expect(topic["similarity"]).to eq(expected_ids.length - idx)
+
+            # blurb is first post cooked
+            expect(topic["blurb"]).to eq(topic.posts.first.cooked)
+          end
+        ensure
+          DiscoursePluginRegistry.unregister_modifier(
+            plugin_instance,
+            :similar_topic_candidate_ids,
+            &blk
+          )
+        end
+      end
+    end
+
     context "with a similar topic" do
       fab!(:post) do
         with_search_indexer_enabled do