From 6da35d8e660f6cbf7935a579cdd184161fdc425e Mon Sep 17 00:00:00 2001
From: Roman Rizzi <roman@discourse.org>
Date: Tue, 10 Dec 2024 15:42:31 -0300
Subject: [PATCH] FIX: Gemini inference client was missing #instance (#1019)

---
 .../vector_representations/gemini.rb           |  2 +-
 lib/inference/gemini_embeddings.rb             |  6 +++++-
 .../vector_representations/gemini_spec.rb      | 18 ++++++++++++++++++
 spec/support/embeddings_generation_stubs.rb    | 10 ++++++++++
 4 files changed, 34 insertions(+), 2 deletions(-)
 create mode 100644 spec/lib/modules/embeddings/vector_representations/gemini_spec.rb

diff --git a/lib/embeddings/vector_representations/gemini.rb b/lib/embeddings/vector_representations/gemini.rb
index a693849d..5ac1cba7 100644
--- a/lib/embeddings/vector_representations/gemini.rb
+++ b/lib/embeddings/vector_representations/gemini.rb
@@ -43,7 +43,7 @@ module DiscourseAi
         end
 
         def vector_from(text, asymetric: false)
-          inference_client.perform!(text).dig(:embedding, :values)
+          inference_client.perform!(text)
         end
 
         # There is no public tokenizer for Gemini, and from the ones we already ship in the plugin
diff --git a/lib/inference/gemini_embeddings.rb b/lib/inference/gemini_embeddings.rb
index 13fb62c5..2684dcd4 100644
--- a/lib/inference/gemini_embeddings.rb
+++ b/lib/inference/gemini_embeddings.rb
@@ -3,6 +3,10 @@
 module ::DiscourseAi
   module Inference
     class GeminiEmbeddings
+      def self.instance
+        new(SiteSetting.ai_gemini_api_key)
+      end
+
       def initialize(api_key, referer = Discourse.base_url)
         @api_key = api_key
         @referer = referer
@@ -21,7 +25,7 @@ module ::DiscourseAi
 
         case response.status
         when 200
-          JSON.parse(response.body, symbolize_names: true)
+          JSON.parse(response.body, symbolize_names: true).dig(:embedding, :values)
         when 429
           # TODO add a AdminDashboard Problem?
         else
diff --git a/spec/lib/modules/embeddings/vector_representations/gemini_spec.rb b/spec/lib/modules/embeddings/vector_representations/gemini_spec.rb
new file mode 100644
index 00000000..80f021f7
--- /dev/null
+++ b/spec/lib/modules/embeddings/vector_representations/gemini_spec.rb
@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+require_relative "vector_rep_shared_examples"
+
+RSpec.describe DiscourseAi::Embeddings::VectorRepresentations::Gemini do
+  subject(:vector_rep) { described_class.new(truncation) }
+
+  let(:truncation) { DiscourseAi::Embeddings::Strategies::Truncation.new }
+  let!(:api_key) { "test-123" }
+
+  before { SiteSetting.ai_gemini_api_key = api_key }
+
+  def stub_vector_mapping(text, expected_embedding)
+    EmbeddingsGenerationStubs.gemini_service(api_key, text, expected_embedding)
+  end
+
+  it_behaves_like "generates and store embedding using with vector representation"
+end
diff --git a/spec/support/embeddings_generation_stubs.rb b/spec/support/embeddings_generation_stubs.rb
index 06ad6352..283ed7d9 100644
--- a/spec/support/embeddings_generation_stubs.rb
+++ b/spec/support/embeddings_generation_stubs.rb
@@ -19,5 +19,15 @@ class EmbeddingsGenerationStubs
         .with(body: JSON.dump({ model: model, input: string }.merge(extra_args)))
         .to_return(status: 200, body: JSON.dump({ data: [{ embedding: embedding }] }))
     end
+
+    def gemini_service(api_key, string, embedding)
+      WebMock
+        .stub_request(
+          :post,
+          "https://generativelanguage.googleapis.com/v1beta/models/embedding-001:embedContent\?key\=#{api_key}",
+        )
+        .with(body: JSON.dump({ content: { parts: [{ text: string }] } }))
+        .to_return(status: 200, body: JSON.dump({ embedding: { values: embedding } }))
+    end
   end
 end