FIX: Open AI embedding shortening is only available for some models (#1080)

Author: romanrizzi

app/models/embedding_definition.rb

@@ -73,6 +73,8 @@ def presets
               tokenizer_class: "DiscourseAi::Tokenizer::MultilingualE5LargeTokenizer",
               provider: HUGGING_FACE,
             },
+            # "text-embedding-3-large" real dimentions are 3072, but we only support up to 2000 in the
+            # indexes, so we downsample to 2000 via API.
             {
               preset_id: "text-embedding-3-large",
               display_name: "OpenAI's text-embedding-3-large",
@@ -198,11 +200,15 @@ def hugging_face_client
   end
 
   def open_ai_client
+    model_name = lookup_custom_param("model_name")
+    can_shorten_dimensions = %w[text-embedding-3-small text-embedding-3-large].include?(model_name)
+    client_dimensions = can_shorten_dimensions ? dimensions : nil
+
     DiscourseAi::Inference::OpenAiEmbeddings.new(
       endpoint_url,
       api_key,
       lookup_custom_param("model_name"),
-      dimensions,
+      client_dimensions,
     )
   end
 

spec/lib/modules/embeddings/vector_spec.rb

@@ -94,13 +94,35 @@ def stub_vector_mapping(text, expected_embedding)
         vdef.lookup_custom_param("model_name"),
         text,
         expected_embedding,
-        extra_args: {
-          dimensions: vdef.dimensions,
-        },
       )
     end
 
     it_behaves_like "generates and store embeddings using a vector definition"
+
+    context "when working with models that support shortening embeddings" do
+      it "passes the dimensions param" do
+        shorter_dimensions = 10
+        vdef.update!(
+          dimensions: shorter_dimensions,
+          provider_params: {
+            model_name: "text-embedding-3-small",
+          },
+        )
+        text = "This is a piece of text"
+        short_expected_embedding = [0.0038493] * shorter_dimensions
+
+        EmbeddingsGenerationStubs.openai_service(
+          vdef.lookup_custom_param("model_name"),
+          text,
+          short_expected_embedding,
+          extra_args: {
+            dimensions: shorter_dimensions,
+          },
+        )
+
+        expect(described_class.new(vdef).vector_from(text)).to eq(short_expected_embedding)
+      end
+    end
   end
 
   context "with hugging_face as the provider" do

spec/requests/embeddings/embeddings_controller_spec.rb

@@ -38,9 +38,6 @@ def stub_embedding(query)
         vector_def.lookup_custom_param("model_name"),
         query,
         embedding,
-        extra_args: {
-          dimensions: vector_def.dimensions,
-        },
       )
     end