FEATURE: configurable embeddings (#1049)

* Use AR model for embeddings features

* endpoints

* Embeddings CRUD UI

* Add presets. Hide a couple more settings

* system specs

* Seed embedding definition from old settings

* Generate search bit index on the fly. cleanup orphaned data

* support for seeded models

* Fix run test for new embedding

* fix selected model not set correctly

Author: romanrizzi

admin/assets/javascripts/discourse/routes/admin-plugins-show-discourse-ai-embeddings-edit.js

@@ -0,0 +1,21 @@
+import DiscourseRoute from "discourse/routes/discourse";
+
+export default class AdminPluginsShowDiscourseAiEmbeddingsEdit extends DiscourseRoute {
+  async model(params) {
+    const allEmbeddings = this.modelFor(
+      "adminPlugins.show.discourse-ai-embeddings"
+    );
+    const id = parseInt(params.id, 10);
+    const record = allEmbeddings.findBy("id", id);
+    record.provider_params = record.provider_params || {};
+    return record;
+  }
+
+  setupController(controller, model) {
+    super.setupController(controller, model);
+    controller.set(
+      "allEmbeddings",
+      this.modelFor("adminPlugins.show.discourse-ai-embeddings")
+    );
+  }
+}

admin/assets/javascripts/discourse/routes/admin-plugins-show-discourse-ai-embeddings-new.js

@@ -0,0 +1,17 @@
+import DiscourseRoute from "discourse/routes/discourse";
+
+export default class AdminPluginsShowDiscourseAiEmbeddingsNew extends DiscourseRoute {
+  async model() {
+    const record = this.store.createRecord("ai-embedding");
+    record.provider_params = {};
+    return record;
+  }
+
+  setupController(controller, model) {
+    super.setupController(controller, model);
+    controller.set(
+      "allEmbeddings",
+      this.modelFor("adminPlugins.show.discourse-ai-embeddings")
+    );
+  }
+}

admin/assets/javascripts/discourse/routes/admin-plugins-show-discourse-ai-embeddings.js

@@ -0,0 +1,7 @@
+import DiscourseRoute from "discourse/routes/discourse";
+
+export default class DiscourseAiAiEmbeddingsRoute extends DiscourseRoute {
+  model() {
+    return this.store.findAll("ai-embedding");
+  }
+}

admin/assets/javascripts/discourse/templates/admin-plugins/show/discourse-ai-embeddings/edit.hbs

@@ -0,0 +1,4 @@
+<AiEmbeddingsListEditor
+  @embeddings={{this.allEmbeddings}}
+  @currentEmbedding={{this.model}}
+/>

admin/assets/javascripts/discourse/templates/admin-plugins/show/discourse-ai-embeddings/index.hbs

@@ -0,0 +1 @@
+<AiEmbeddingsListEditor @embeddings={{this.model}} />

admin/assets/javascripts/discourse/templates/admin-plugins/show/discourse-ai-embeddings/new.hbs

@@ -0,0 +1,4 @@
+<AiEmbeddingsListEditor
+  @embeddings={{this.allEmbeddings}}
+  @currentEmbedding={{this.model}}
+/>

app/controllers/discourse_ai/admin/ai_embeddings_controller.rb

@@ -0,0 +1,130 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module Admin
+    class AiEmbeddingsController < ::Admin::AdminController
+      requires_plugin ::DiscourseAi::PLUGIN_NAME
+
+      def index
+        embedding_defs = EmbeddingDefinition.all.order(:display_name)
+
+        render json: {
+                 ai_embeddings:
+                   ActiveModel::ArraySerializer.new(
+                     embedding_defs,
+                     each_serializer: AiEmbeddingDefinitionSerializer,
+                     root: false,
+                   ).as_json,
+                 meta: {
+                   provider_params: EmbeddingDefinition.provider_params,
+                   providers: EmbeddingDefinition.provider_names,
+                   distance_functions: EmbeddingDefinition.distance_functions,
+                   tokenizers:
+                     EmbeddingDefinition.tokenizer_names.map { |tn|
+                       { id: tn, name: tn.split("::").last }
+                     },
+                   presets: EmbeddingDefinition.presets,
+                 },
+               }
+      end
+
+      def new
+      end
+
+      def edit
+        embedding_def = EmbeddingDefinition.find(params[:id])
+        render json: AiEmbeddingDefinitionSerializer.new(embedding_def)
+      end
+
+      def create
+        embedding_def = EmbeddingDefinition.new(ai_embeddings_params)
+
+        if embedding_def.save
+          render json: AiEmbeddingDefinitionSerializer.new(embedding_def), status: :created
+        else
+          render_json_error embedding_def
+        end
+      end
+
+      def update
+        embedding_def = EmbeddingDefinition.find(params[:id])
+
+        if embedding_def.seeded?
+          return(
+            render_json_error(I18n.t("discourse_ai.embeddings.cannot_edit_builtin"), status: 403)
+          )
+        end
+
+        if embedding_def.update(ai_embeddings_params.except(:dimensions))
+          render json: AiEmbeddingDefinitionSerializer.new(embedding_def)
+        else
+          render_json_error embedding_def
+        end
+      end
+
+      def destroy
+        embedding_def = EmbeddingDefinition.find(params[:id])
+
+        if embedding_def.seeded?
+          return(
+            render_json_error(I18n.t("discourse_ai.embeddings.cannot_edit_builtin"), status: 403)
+          )
+        end
+
+        if embedding_def.id == SiteSetting.ai_embeddings_selected_model.to_i
+          return render_json_error(I18n.t("discourse_ai.embeddings.delete_failed"), status: 409)
+        end
+
+        if embedding_def.destroy
+          head :no_content
+        else
+          render_json_error embedding_def
+        end
+      end
+
+      def test
+        RateLimiter.new(
+          current_user,
+          "ai_embeddings_test_#{current_user.id}",
+          3,
+          1.minute,
+        ).performed!
+
+        embedding_def = EmbeddingDefinition.new(ai_embeddings_params)
+        DiscourseAi::Embeddings::Vector.new(embedding_def).vector_from("this is a test")
+
+        render json: { success: true }
+      rescue Net::HTTPBadResponse => e
+        render json: { success: false, error: e.message }
+      end
+
+      private
+
+      def ai_embeddings_params
+        permitted =
+          params.require(:ai_embedding).permit(
+            :display_name,
+            :dimensions,
+            :max_sequence_length,
+            :pg_function,
+            :provider,
+            :url,
+            :api_key,
+            :tokenizer_class,
+          )
+
+        extra_field_names = EmbeddingDefinition.provider_params.dig(permitted[:provider]&.to_sym)
+        if extra_field_names.present?
+          received_prov_params =
+            params.dig(:ai_embedding, :provider_params)&.slice(*extra_field_names.keys)
+
+          if received_prov_params.present?
+            permitted[:provider_params] = received_prov_params.permit!
+          end
+        end
+
+        permitted
+      end
+    end
+  end
+end

app/jobs/regular/digest_rag_upload.rb

@@ -18,7 +18,7 @@ def execute(args)
       target = target_type.constantize.find_by(id: target_id)
       return if !target
 
-      vector_rep = DiscourseAi::Embeddings::VectorRepresentations::Base.current_representation
+      vector_rep = DiscourseAi::Embeddings::Vector.instance
 
       tokenizer = vector_rep.tokenizer
       chunk_tokens = target.rag_chunk_tokens

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