Generate search bit index on the fly. cleanup orphaned data

Author: romanrizzi

app/controllers/discourse_ai/admin/ai_embeddings_controller.rb

@@ -49,7 +49,7 @@ def create
       def update
         embedding_def = EmbeddingDefinition.find(params[:id])
 
-        if embedding_def.update(ai_embeddings_params)
+        if embedding_def.update(ai_embeddings_params.except(:dimensions))
           render json: AiEmbeddingDefinitionSerializer.new(embedding_def)
         else
           render_json_error embedding_def

app/jobs/regular/manage_embedding_def_search_index.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module ::Jobs
+  class ManageEmbeddingDefSearchIndex < ::Jobs::Base
+    def execute(args)
+      embedding_def = EmbeddingDefinition.find_by(id: args[:id])
+      return if embedding_def.nil?
+      return if DiscourseAi::Embeddings::Schema.correctly_indexed?(embedding_def)
+
+      DiscourseAi::Embeddings::Schema.prepare_search_indexes(embedding_def)
+    end
+  end
+end

app/jobs/scheduled/remove_orphaned_embeddings.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Jobs
+  class RemoveOrphanedEmbeddings < ::Jobs::Scheduled
+    every 1.week
+
+    def execute(_args)
+      DiscourseAi::Embeddings::Schema.remove_orphaned_data
+    end
+  end
+end

app/models/embedding_definition.rb

@@ -121,6 +121,12 @@ def presets
   validates :tokenizer_class, presence: true, inclusion: tokenizer_names
   validates_presence_of :url, :api_key, :dimensions, :max_sequence_length, :pg_function
 
+  after_create :create_indexes
+
+  def create_indexes
+    Jobs.enqueue(:manage_embedding_def_search_index, id: self.id)
+  end
+
   def tokenizer
     tokenizer_class.constantize
   end

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

@@ -10,10 +10,12 @@ import { later } from "@ember/runloop";
 import { service } from "@ember/service";
 import BackButton from "discourse/components/back-button";
 import DButton from "discourse/components/d-button";
+import icon from "discourse/helpers/d-icon";
 import { popupAjaxError } from "discourse/lib/ajax-error";
-import icon from "discourse-common/helpers/d-icon";
 import { i18n } from "discourse-i18n";
 import ComboBox from "select-kit/components/combo-box";
+import DTooltip from "float-kit/components/d-tooltip";
+import not from "truth-helpers/helpers/not";
 
 export default class AiEmbeddingEditor extends Component {
   @service toasts;
@@ -276,7 +278,16 @@ export default class AiEmbeddingEditor extends Component {
             lang="en"
             @value={{this.editingModel.dimensions}}
             required="true"
+            disabled={{not this.editingModel.isNew}}
           />
+          {{#if this.editingModel.isNew}}
+            <DTooltip
+              @icon="circle-exclamation"
+              @content={{i18n
+                "discourse_ai.embeddings.hints.dimensions_warning"
+              }}
+            />
+          {{/if}}
         </div>
 
         <div class="control-group">

config/locales/client.en.yml

@@ -519,6 +519,8 @@ en:
           running: "Running test..."
           success: "Success!"
           failure: "Attempting to generate an embedding resulted in: %{error}"
+        hints:
+          dimensions_warning: "Once saved, this value can't be changed."
 
         display_name: "Name"
         provider: "Provider"

lib/embeddings/schema.rb

@@ -12,21 +12,80 @@ class Schema
       POSTS_TABLE = "ai_posts_embeddings"
       RAG_DOCS_TABLE = "ai_document_fragments_embeddings"
 
+      EMBEDDING_TARGETS = %w[topics posts document_fragments]
+      EMBEDDING_TABLES = [TOPICS_TABLE, POSTS_TABLE, RAG_DOCS_TABLE]
+
       MissingEmbeddingError = Class.new(StandardError)
 
-      def self.for(target_klass)
-        vector_def = EmbeddingDefinition.find_by(id: SiteSetting.ai_embeddings_selected_model)
-        raise "Invalid embeddings selected model" if vector_def.nil?
-
-        case target_klass&.name
-        when "Topic"
-          new(TOPICS_TABLE, "topic_id", vector_def)
-        when "Post"
-          new(POSTS_TABLE, "post_id", vector_def)
-        when "RagDocumentFragment"
-          new(RAG_DOCS_TABLE, "rag_document_fragment_id", vector_def)
-        else
-          raise ArgumentError, "Invalid target type for embeddings"
+      class << self
+        def for(target_klass)
+          vector_def = EmbeddingDefinition.find_by(id: SiteSetting.ai_embeddings_selected_model)
+          raise "Invalid embeddings selected model" if vector_def.nil?
+
+          case target_klass&.name
+          when "Topic"
+            new(TOPICS_TABLE, "topic_id", vector_def)
+          when "Post"
+            new(POSTS_TABLE, "post_id", vector_def)
+          when "RagDocumentFragment"
+            new(RAG_DOCS_TABLE, "rag_document_fragment_id", vector_def)
+          else
+            raise ArgumentError, "Invalid target type for embeddings"
+          end
+        end
+
+        def search_index_name(table, def_id)
+          "ai_#{table}_embeddings_#{def_id}_1_search_bit"
+        end
+
+        def prepare_search_indexes(vector_def)
+          EMBEDDING_TARGETS.each { |target| DB.exec <<~SQL }
+              CREATE INDEX IF NOT EXISTS #{search_index_name(target, vector_def.id)} ON ai_#{target}_embeddings
+              USING hnsw ((binary_quantize(embeddings)::bit(#{vector_def.dimensions})) bit_hamming_ops)
+              WHERE model_id = #{vector_def.id} AND strategy_id = 1;
+            SQL
+        end
+
+        def correctly_indexed?(vector_def)
+          index_names = EMBEDDING_TARGETS.map { |t| search_index_name(t, vector_def.id) }
+          indexdefs =
+            DB.query_single(
+              "SELECT indexdef FROM pg_indexes WHERE indexname IN (:names)",
+              names: index_names,
+            )
+
+          return false if indexdefs.length < index_names.length
+
+          indexdefs.all? do |defs|
+            defs.include? "(binary_quantize(embeddings))::bit(#{vector_def.dimensions})"
+          end
+        end
+
+        def remove_orphaned_data
+          removed_defs_ids =
+            DB.query_single(
+              "SELECT DISTINCT(model_id) FROM #{TOPICS_TABLE} te LEFT JOIN embedding_definitions ed ON te.model_id = ed.id WHERE ed.id IS NULL",
+            )
+
+          EMBEDDING_TABLES.each do |t|
+            DB.exec(
+              "DELETE FROM #{t} WHERE model_id IN (:removed_defs)",
+              removed_defs: removed_defs_ids,
+            )
+          end
+
+          drop_index_statement =
+            EMBEDDING_TARGETS
+              .reduce([]) do |memo, et|
+                removed_defs_ids.each do |rdi|
+                  memo << "DROP INDEX IF EXISTS #{search_index_name(et, rdi)};"
+                end
+
+                memo
+              end
+              .join("\n")
+
+          DB.exec(drop_index_statement)
         end
       end
 

spec/jobs/regular/manage_embedding_def_search_index_spec.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+RSpec.describe Jobs::ManageEmbeddingDefSearchIndex do
+  fab!(:embedding_definition)
+
+  describe "#execute" do
+    context "when there is no embedding def" do
+      it "does nothing" do
+        invalid_id = 999_999_999
+
+        subject.execute(id: invalid_id)
+
+        expect(
+          DiscourseAi::Embeddings::Schema.correctly_indexed?(
+            EmbeddingDefinition.new(id: invalid_id),
+          ),
+        ).to eq(false)
+      end
+    end
+
+    context "when the embedding def is fresh" do
+      it "creates the indexes" do
+        subject.execute(id: embedding_definition.id)
+
+        expect(DiscourseAi::Embeddings::Schema.correctly_indexed?(embedding_definition)).to eq(true)
+      end
+
+      it "creates them only once" do
+        subject.execute(id: embedding_definition.id)
+        subject.execute(id: embedding_definition.id)
+
+        expect(DiscourseAi::Embeddings::Schema.correctly_indexed?(embedding_definition)).to eq(true)
+      end
+
+      context "when one of the idxs is missing" do
+        it "automatically recovers by creating it" do
+          DB.exec <<~SQL
+            CREATE INDEX IF NOT EXISTS ai_topics_embeddings_#{embedding_definition.id}_1_search_bit ON ai_topics_embeddings
+            USING hnsw ((binary_quantize(embeddings)::bit(#{embedding_definition.dimensions})) bit_hamming_ops)
+            WHERE model_id = #{embedding_definition.id} AND strategy_id = 1;
+          SQL
+
+          expect(DiscourseAi::Embeddings::Schema.correctly_indexed?(embedding_definition)).to eq(
+            false,
+          )
+
+          subject.execute(id: embedding_definition.id)
+
+          expect(DiscourseAi::Embeddings::Schema.correctly_indexed?(embedding_definition)).to eq(
+            true,
+          )
+        end
+      end
+    end
+  end
+end

spec/jobs/scheduled/remove_orphaned_embeddings_spec.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+RSpec.describe Jobs::RemoveOrphanedEmbeddings do
+  describe "#execute" do
+    fab!(:embedding_definition)
+    fab!(:embedding_definition_2) { Fabricate(:embedding_definition) }
+    fab!(:topic)
+    fab!(:post)
+
+    before do
+      DiscourseAi::Embeddings::Schema.prepare_search_indexes(embedding_definition)
+      DiscourseAi::Embeddings::Schema.prepare_search_indexes(embedding_definition_2)
+
+      # Seed embeddings. One of each def x target classes.
+      [embedding_definition, embedding_definition_2].each do |edef|
+        SiteSetting.ai_embeddings_selected_model = edef.id
+
+        [topic, post].each do |target|
+          schema = DiscourseAi::Embeddings::Schema.for(target.class)
+          schema.store(target, [1] * edef.dimensions, "test")
+        end
+      end
+
+      embedding_definition.destroy!
+    end
+
+    def find_all_embeddings_of(target, table, target_column)
+      DB.query_single("SELECT model_id FROM #{table} WHERE #{target_column} = #{target.id}")
+    end
+
+    it "delete embeddings without an existing embedding definition" do
+      expect(find_all_embeddings_of(post, "ai_posts_embeddings", "post_id")).to contain_exactly(
+        embedding_definition.id,
+        embedding_definition_2.id,
+      )
+      expect(find_all_embeddings_of(topic, "ai_topics_embeddings", "topic_id")).to contain_exactly(
+        embedding_definition.id,
+        embedding_definition_2.id,
+      )
+
+      subject.execute({})
+
+      expect(find_all_embeddings_of(topic, "ai_topics_embeddings", "topic_id")).to contain_exactly(
+        embedding_definition_2.id,
+      )
+      expect(find_all_embeddings_of(post, "ai_posts_embeddings", "post_id")).to contain_exactly(
+        embedding_definition_2.id,
+      )
+    end
+
+    it "deletes orphaned indexes" do
+      expect(DiscourseAi::Embeddings::Schema.correctly_indexed?(embedding_definition)).to eq(true)
+      expect(DiscourseAi::Embeddings::Schema.correctly_indexed?(embedding_definition_2)).to eq(true)
+
+      subject.execute({})
+
+      index_names =
+        DiscourseAi::Embeddings::Schema::EMBEDDING_TARGETS.map do |t|
+          "ai_#{t}_embeddings_#{embedding_definition.id}_1_search_bit"
+        end
+      indexnames =
+        DB.query_single(
+          "SELECT indexname FROM pg_indexes WHERE indexname IN (:names)",
+          names: index_names,
+        )
+
+      expect(indexnames).to be_empty
+      expect(DiscourseAi::Embeddings::Schema.correctly_indexed?(embedding_definition_2)).to eq(true)
+    end
+  end
+end

spec/requests/admin/ai_embeddings_controller_spec.rb

@@ -95,6 +95,20 @@
 
         expect(response.status).to eq(404)
       end
+
+      it "doesn't allow dimenstions to be updated" do
+        new_dimensions = 200
+
+        put "/admin/plugins/discourse-ai/ai-embeddings/#{embedding_definition.id}.json",
+            params: {
+              ai_embedding: {
+                dimensions: new_dimensions,
+              },
+            }
+
+        expect(response.status).to eq(200)
+        expect(embedding_definition.reload.dimensions).not_to eq(new_dimensions)
+      end
     end
 
     context "with invalid update params" do