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

This commit is contained in:
Roman Rizzi 2025-01-21 17:50:40 -03:00 committed by GitHub
parent 3b66fb3e87
commit a5e5ae72a8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 32 additions and 7 deletions

View File

@ -73,6 +73,8 @@ class EmbeddingDefinition < ActiveRecord::Base
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 @@ class EmbeddingDefinition < ActiveRecord::Base
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

View File

@ -94,13 +94,35 @@ RSpec.describe DiscourseAi::Embeddings::Vector do
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

View File

@ -38,9 +38,6 @@ describe DiscourseAi::Embeddings::EmbeddingsController do
vector_def.lookup_custom_param("model_name"),
query,
embedding,
extra_args: {
dimensions: vector_def.dimensions,
},
)
end