diff --git a/app/models/embedding_definition.rb b/app/models/embedding_definition.rb index 8f0bd5c2f..0c7cbda4c 100644 --- a/app/models/embedding_definition.rb +++ b/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 diff --git a/spec/lib/modules/embeddings/vector_spec.rb b/spec/lib/modules/embeddings/vector_spec.rb index c6e8b8c58..b69073c38 100644 --- a/spec/lib/modules/embeddings/vector_spec.rb +++ b/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 diff --git a/spec/requests/embeddings/embeddings_controller_spec.rb b/spec/requests/embeddings/embeddings_controller_spec.rb index 9ad04d5b4..0408aa37d 100644 --- a/spec/requests/embeddings/embeddings_controller_spec.rb +++ b/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