FEATURE: Support for Gemini Embeddings (#382)

config/settings.yml

@@ -217,6 +217,7 @@ discourse_ai:
       - text-embedding-ada-002
       - multilingual-e5-large
       - bge-large-en
+      - gemini
   ai_embeddings_generate_for_pms: false
   ai_embeddings_semantic_related_topics_enabled:
     default: false

db/migrate/20231227223301_create_gemini_topic_embeddings_table.rb

@@ -0,0 +1,16 @@
+# frozen_string_literal: true
+
+class CreateGeminiTopicEmbeddingsTable < ActiveRecord::Migration[7.0]
+  def change
+    create_table :ai_topic_embeddings_5_1, id: false do |t|
+      t.integer :topic_id, null: false
+      t.integer :model_version, null: false
+      t.integer :strategy_version, null: false
+      t.text :digest, null: false
+      t.column :embeddings, "vector(768)", null: false
+      t.timestamps
+
+      t.index :topic_id, unique: true
+    end
+  end
+end

lib/embeddings/vector_representations/base.rb

@@ -10,6 +10,7 @@ module DiscourseAi
           [
             DiscourseAi::Embeddings::VectorRepresentations::AllMpnetBaseV2,
             DiscourseAi::Embeddings::VectorRepresentations::BgeLargeEn,
+            DiscourseAi::Embeddings::VectorRepresentations::Gemini,
             DiscourseAi::Embeddings::VectorRepresentations::MultilingualE5Large,
             DiscourseAi::Embeddings::VectorRepresentations::TextEmbeddingAda002,
           ].map { _1.new(strategy) }.find { _1.name == SiteSetting.ai_embeddings_model }

lib/embeddings/vector_representations/gemini.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module DiscourseAi
+  module Embeddings
+    module VectorRepresentations
+      class Gemini < Base
+        def id
+          5
+        end
+
+        def version
+          1
+        end
+
+        def name
+          "gemini"
+        end
+
+        def dimensions
+          768
+        end
+
+        def max_sequence_length
+          2048
+        end
+
+        def pg_function
+          "<=>"
+        end
+
+        def pg_index_type
+          "vector_cosine_ops"
+        end
+
+        def vector_from(text)
+          response = DiscourseAi::Inference::GeminiEmbeddings.perform!(text)
+          response[:embedding][:values]
+        end
+
+        # There is no public tokenizer for Gemini, and from the ones we already ship in the plugin
+        # OpenAI gets the closest results. Gemini Tokenizer results in ~10% less tokens, so it's safe
+        # to use OpenAI tokenizer since it will overestimate the number of tokens.
+        def tokenizer
+          DiscourseAi::Tokenizer::OpenAiTokenizer
+        end
+      end
+    end
+  end
+end

lib/inference/gemini_embeddings.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module ::DiscourseAi
+  module Inference
+    class GeminiEmbeddings
+      def self.perform!(content)
+        headers = { "Referer" => Discourse.base_url, "Content-Type" => "application/json" }
+
+        url =
+          "https://generativelanguage.googleapis.com/v1beta/models/embedding-001:embedContent\?key\=#{SiteSetting.ai_gemini_api_key}"
+
+        body = { content: { parts: [{ text: content }] } }
+
+        response = Faraday.post(url, body.to_json, headers)
+
+        raise Net::HTTPBadResponse if ![200].include?(response.status)
+
+        JSON.parse(response.body, symbolize_names: true)
+      end
+    end
+  end
+end