Use AR model for embeddings features

Author: romanrizzi

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/models/embedding_definition.rb

@@ -0,0 +1,139 @@
+# frozen_string_literal: true
+
+class EmbeddingDefinition < ActiveRecord::Base
+  CLOUDFLARE = "cloudflare"
+  DISCOURSE = "discourse"
+  HUGGING_FACE = "hugging_face"
+  OPEN_AI = "open_ai"
+  GEMINI = "gemini"
+
+  class << self
+    def provider_names
+      [CLOUDFLARE, DISCOURSE, HUGGING_FACE, OPEN_AI, GEMINI]
+    end
+
+    def tokenizer_names
+      [
+        DiscourseAi::Tokenizer::AllMpnetBaseV2Tokenizer,
+        DiscourseAi::Tokenizer::BgeLargeEnTokenizer,
+        DiscourseAi::Tokenizer::BgeM3Tokenizer,
+        DiscourseAi::Tokenizer::OpenAiTokenizer,
+        DiscourseAi::Tokenizer::MultilingualE5LargeTokenizer,
+        DiscourseAi::Tokenizer::OpenAiTokenizer,
+      ].map(&:name)
+    end
+
+    def self.provider_params
+      { discourse: { model_name: :text }, open_ai: { model_name: :text } }
+    end
+  end
+
+  validates :provider, presence: true, inclusion: provider_names
+  validates :display_name, presence: true, length: { maximum: 100 }
+  validates :tokenizer_class, presence: true, inclusion: tokenizer_names
+
+  def tokenizer
+    tokenizer_class.constantize
+  end
+
+  def inference_client
+    case provider
+    when CLOUDFLARE
+      cloudflare_client
+    when DISCOURSE
+      discourse_client
+    when HUGGING_FACE
+      hugging_face_client
+    when OPEN_AI
+      open_ai_client
+    when GEMINI
+      gemini_client
+    else
+      raise "Uknown embeddings provider"
+    end
+  end
+
+  def lookup_custom_param(key)
+    provider_params&.dig(key)
+  end
+
+  def endpoint_url
+    return url if !url.starts_with?("srv://")
+
+    service = DiscourseAi::Utils::DnsSrv.lookup(url)
+    "https://#{service.target}:#{service.port}"
+  end
+
+  def prepare_query_text(text, asymetric: false)
+    strategy.prepare_query_text(text, self, asymetric: asymetric)
+  end
+
+  def prepare_target_text(target)
+    strategy.prepare_target_text(target, self)
+  end
+
+  def strategy_id
+    strategy.id
+  end
+
+  def strategy_version
+    strategy.version
+  end
+
+  private
+
+  def strategy
+    @strategy ||= DiscourseAi::Embeddings::Strategies::Truncation.new
+  end
+
+  def cloudflare_client
+    DiscourseAi::Inference::CloudflareWorkersAi.new(endpoint_url, api_key)
+  end
+
+  def discourse_client
+    client_url = endpoint_url
+    client_url = "#{client_url}/api/v1/classify" if url.starts_with?("srv://")
+
+    DiscourseAi::Inference::DiscourseClassifier.new(
+      client_url,
+      api_key,
+      lookup_custom_param("model_name"),
+    )
+  end
+
+  def hugging_face_client
+    DiscourseAi::Inference::HuggingFaceTextEmbeddings.new(endpoint_url, api_key)
+  end
+
+  def open_ai_client
+    DiscourseAi::Inference::OpenAiEmbeddings.new(
+      endpoint_url,
+      api_key,
+      lookup_custom_param("model_name"),
+      dimensions,
+    )
+  end
+
+  def gemini_client
+    DiscourseAi::Inference::GeminiEmbeddings.new(endpoint_url, api_key)
+  end
+end
+
+# == Schema Information
+#
+# Table name: embedding_definitions
+#
+#  id                  :bigint           not null, primary key
+#  display_name        :string           not null
+#  dimensions          :integer          not null
+#  max_sequence_length :integer          not null
+#  version             :integer          default(1), not null
+#  pg_function         :string           not null
+#  provider            :string           not null
+#  tokenizer_class     :string           not null
+#  url                 :string           not null
+#  api_key             :string
+#  provider_params     :jsonb
+#  created_at          :datetime         not null
+#  updated_at          :datetime         not null
+#

config/locales/server.en.yml

@@ -68,14 +68,11 @@ en:
     ai_auto_image_caption_allowed_groups: "Users on these groups can toggle automatic image captioning."
 
     ai_embeddings_enabled: "Enable the embeddings module."
-    ai_embeddings_discourse_service_api_endpoint: "URL where the API is running for the embeddings module"
-    ai_embeddings_discourse_service_api_key: "API key for the embeddings API"
-    ai_embeddings_model: "Use all-mpnet-base-v2 for local and fast inference in english, text-embedding-ada-002 to use OpenAI API (need API key) and multilingual-e5-large for local multilingual embeddings"
+    ai_embeddings_selected_model: "Use the selected model for generating embeddings."
     ai_embeddings_generate_for_pms: "Generate embeddings for personal messages."
     ai_embeddings_semantic_related_topics_enabled: "Use Semantic Search for related topics."
     ai_embeddings_semantic_related_topics: "Maximum number of topics to show in related topic section."
     ai_embeddings_backfill_batch_size: "Number of embeddings to backfill every 15 minutes."
-    ai_embeddings_pg_connection_string: "PostgreSQL connection string for the embeddings module. Needs pgvector extension enabled and a series of tables created. See docs for more info."
     ai_embeddings_semantic_search_enabled: "Enable full-page semantic search."
     ai_embeddings_semantic_quick_search_enabled: "Enable semantic search option in search menu popup."
     ai_embeddings_semantic_related_include_closed_topics: "Include closed topics in semantic search results"
@@ -435,11 +432,7 @@ en:
     embeddings:
       configuration:
         disable_embeddings: "You have to disable 'ai embeddings enabled' first."
-        choose_model: "Set 'ai embeddings model' first."
-        model_unreachable: "We failed to generate a test embedding with this model. Check your settings are correct."
-        hint:
-          one: "Make sure the `%{settings}` setting was configured."
-          other: "Make sure the settings of the provider you want were configured. Options are: %{settings}"
+        choose_model: "Set 'ai embeddings selected model' first."
 
     llm_models:
       missing_provider_param: "%{param} can't be blank"

config/settings.yml

@@ -158,27 +158,12 @@ discourse_ai:
     default: false
     client: true
     validator: "DiscourseAi::Configuration::EmbeddingsModuleValidator"
-  ai_embeddings_discourse_service_api_endpoint: ""
-  ai_embeddings_discourse_service_api_endpoint_srv:
-    default: ""
-    hidden: true
-  ai_embeddings_discourse_service_api_key:
-    default: ""
-    secret: true
-  ai_embeddings_model:
+  ai_embeddings_selected_model:
     type: enum
-    default: "bge-large-en"
+    default: ""
     allow_any: false
-    choices:
-      - all-mpnet-base-v2
-      - text-embedding-ada-002
-      - text-embedding-3-small
-      - text-embedding-3-large
-      - multilingual-e5-large
-      - bge-large-en
-      - gemini
-      - bge-m3
-    validator: "DiscourseAi::Configuration::EmbeddingsModelValidator"
+    enum: "DiscourseAi::Configuration::EmbeddingDefsEnumerator"
+    validator: "DiscourseAi::Configuration::EmbeddingDefsValidator"
   ai_embeddings_per_post_enabled:
     default: false
     hidden: true
@@ -191,9 +176,6 @@ discourse_ai:
   ai_embeddings_backfill_batch_size:
     default: 250
     hidden: true
-  ai_embeddings_pg_connection_string:
-    default: ""
-    hidden: true
   ai_embeddings_semantic_search_enabled:
     default: false
     client: true
@@ -213,6 +195,35 @@ discourse_ai:
     default: false
     client: true
     hidden: true
+
+  ai_embeddings_discourse_service_api_endpoint: 
+    default: ""
+    hidden: true
+  ai_embeddings_discourse_service_api_endpoint_srv:
+    default: ""
+    hidden: true
+  ai_embeddings_discourse_service_api_key:
+    hidden: true
+    default: ""
+    secret: true
+  ai_embeddings_model:
+    hidden: true
+    type: enum
+    default: "bge-large-en"
+    allow_any: false
+    choices:
+      - all-mpnet-base-v2
+      - text-embedding-ada-002
+      - text-embedding-3-small
+      - text-embedding-3-large
+      - multilingual-e5-large
+      - bge-large-en
+      - gemini
+      - bge-m3
+  ai_embeddings_pg_connection_string:
+    default: ""
+    hidden: true
+
   ai_summarization_enabled:
     default: false
     client: true

db/migrate/20241217164540_create_embedding_definitions.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+class CreateEmbeddingDefinitions < ActiveRecord::Migration[7.2]
+  def change
+    create_table :embedding_definitions do |t|
+      t.string :display_name, null: false
+      t.integer :dimensions, null: false
+      t.integer :max_sequence_length, null: false
+      t.integer :version, null: false, default: 1
+      t.string :pg_function, null: false
+      t.string :provider, null: false
+      t.string :tokenizer_class, null: false
+      t.string :url, null: false
+      t.string :api_key
+      t.jsonb :provider_params
+      t.timestamps
+    end
+  end
+end

db/migrate/20241230153300_new_embeddings_tables.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+class NewEmbeddingsTables < ActiveRecord::Migration[7.2]
+  def up
+    # We need "model_id" to be a bigint now that it's no longer a hardcoded ID.
+
+    create_table :ai_topics_embeddings, id: false do |t|
+      t.bigint :topic_id, null: false
+      t.bigint :model_id, null: false
+      t.integer :model_version, null: false
+      t.integer :strategy_id, null: false
+      t.integer :strategy_version, null: false
+      t.text :digest, null: false
+      t.column :embeddings, "halfvec", null: false
+      t.timestamps
+
+      t.index %i[model_id strategy_id topic_id],
+              unique: true,
+              name: "index_ai_topics_embeddings_on_model_strategy_topic"
+    end
+
+    create_table :ai_posts_embeddings, id: false do |t|
+      t.bigint :post_id, null: false
+      t.bigint :model_id, null: false
+      t.integer :model_version, null: false
+      t.integer :strategy_id, null: false
+      t.integer :strategy_version, null: false
+      t.text :digest, null: false
+      t.column :embeddings, "halfvec", null: false
+      t.timestamps
+
+      t.index %i[model_id strategy_id post_id],
+              unique: true,
+              name: "index_ai_posts_embeddings_on_model_strategy_post"
+    end
+
+    create_table :ai_document_fragments_embeddings, id: false do |t|
+      t.bigint :rag_document_fragment_id, null: false
+      t.bigint :model_id, null: false
+      t.integer :model_version, null: false
+      t.integer :strategy_id, null: false
+      t.integer :strategy_version, null: false
+      t.text :digest, null: false
+      t.column :embeddings, "halfvec", null: false
+      t.timestamps
+
+      t.index %i[model_id strategy_id rag_document_fragment_id],
+              unique: true,
+              name: "index_ai_fragments_embeddings_on_model_strategy_fragment"
+    end
+
+    # Copy data from old tables to new tables
+    execute <<-SQL
+      INSERT INTO ai_topics_embeddings (topic_id, model_id, model_version, strategy_id, strategy_version, digest, embeddings, created_at, updated_at)
+      SELECT * FROM ai_topic_embeddings;
+
+      INSERT INTO ai_posts_embeddings (post_id, model_id, model_version, strategy_id, strategy_version, digest, embeddings, created_at, updated_at)
+      SELECT * FROM ai_post_embeddings;
+
+      INSERT INTO ai_document_fragments_embeddings (rag_document_fragment_id, model_id, model_version, strategy_id, strategy_version, digest, embeddings, created_at, updated_at)
+      SELECT * FROM ai_document_fragment_embeddings;
+    SQL
+  end
+
+  def down
+    raise ActiveRecord::IrreversibleMigration
+  end
+end

lib/ai_bot/entry_point.rb

@@ -196,7 +196,7 @@ def inject_into(plugin)
         )
 
         plugin.on(:site_setting_changed) do |name, old_value, new_value|
-          if name == :ai_embeddings_model && SiteSetting.ai_embeddings_enabled? &&
+          if name == :ai_embeddings_selected_model && SiteSetting.ai_embeddings_enabled? &&
                new_value != old_value
             RagDocumentFragment.delete_all
             UploadReference

lib/ai_bot/personas/persona.rb

@@ -327,7 +327,7 @@ def rag_fragments_prompt(conversation_context, llm:, user:)
               rag_conversation_chunks
             end
 
-          schema = DiscourseAi::Embeddings::Schema.for(RagDocumentFragment, vector_def: vector.vdef)
+          schema = DiscourseAi::Embeddings::Schema.for(RagDocumentFragment)
 
           candidate_fragment_ids =
             schema

lib/ai_helper/semantic_categorizer.rb

@@ -93,7 +93,7 @@ def tags
 
       def nearest_neighbors(limit: 100)
         vector = DiscourseAi::Embeddings::Vector.instance
-        schema = DiscourseAi::Embeddings::Schema.for(Topic, vector_def: vector.vdef)
+        schema = DiscourseAi::Embeddings::Schema.for(Topic)
 
         raw_vector = vector.vector_from(@text)
 

lib/configuration/embedding_defs_enumerator.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+require "enum_site_setting"
+
+module DiscourseAi
+  module Configuration
+    class EmbeddingDefsEnumerator < ::EnumSiteSetting
+      def self.valid_value?(val)
+        true
+      end
+
+      def self.values
+        DB.query_hash(<<~SQL).map(&:symbolize_keys)
+          SELECT display_name AS name, id AS value
+          FROM embedding_definitions
+        SQL
+      end
+    end
+  end
+end