From 3be76ebd7a83a55d817daded682e42e79d3d1fb9 Mon Sep 17 00:00:00 2001
From: Rafael dos Santos Silva <xfalcox@gmail.com>
Date: Thu, 11 Jan 2024 14:16:25 -0300
Subject: [PATCH] FEATURE: Move the default embeddings model to bge-large-en
 (#417)

---
 config/settings.yml                                          | 2 +-
 spec/lib/modules/ai_bot/tools/search_spec.rb                 | 2 +-
 spec/lib/modules/embeddings/jobs/generate_embeddings_spec.rb | 2 +-
 spec/lib/modules/embeddings/semantic_search_spec.rb          | 2 +-
 spec/lib/modules/embeddings/semantic_topic_query_spec.rb     | 2 +-
 spec/support/embeddings_generation_stubs.rb                  | 1 +
 6 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/config/settings.yml b/config/settings.yml
index 4409503d..75c135d7 100644
--- a/config/settings.yml
+++ b/config/settings.yml
@@ -231,7 +231,7 @@ discourse_ai:
   ai_embeddings_model:
     type: enum
     list_type: compact
-    default: "all-mpnet-base-v2"
+    default: "bge-large-en"
     allow_any: false
     choices:
       - all-mpnet-base-v2
diff --git a/spec/lib/modules/ai_bot/tools/search_spec.rb b/spec/lib/modules/ai_bot/tools/search_spec.rb
index a5136ea4..9f0b8e7c 100644
--- a/spec/lib/modules/ai_bot/tools/search_spec.rb
+++ b/spec/lib/modules/ai_bot/tools/search_spec.rb
@@ -80,7 +80,7 @@ RSpec.describe DiscourseAi::AiBot::Tools::Search do
         post1 = Fabricate(:post, topic: topic_with_tags)
         search = described_class.new({ search_query: "hello world, sam", status: "public" })
 
-        DiscourseAi::Embeddings::VectorRepresentations::AllMpnetBaseV2
+        DiscourseAi::Embeddings::VectorRepresentations::BgeLargeEn
           .any_instance
           .expects(:asymmetric_topics_similarity_search)
           .returns([post1.topic_id])
diff --git a/spec/lib/modules/embeddings/jobs/generate_embeddings_spec.rb b/spec/lib/modules/embeddings/jobs/generate_embeddings_spec.rb
index d8758289..78aee2ae 100644
--- a/spec/lib/modules/embeddings/jobs/generate_embeddings_spec.rb
+++ b/spec/lib/modules/embeddings/jobs/generate_embeddings_spec.rb
@@ -7,7 +7,7 @@ RSpec.describe Jobs::GenerateEmbeddings do
     before do
       SiteSetting.ai_embeddings_discourse_service_api_endpoint = "http://test.com"
       SiteSetting.ai_embeddings_enabled = true
-      SiteSetting.ai_embeddings_model = "all-mpnet-base-v2"
+      SiteSetting.ai_embeddings_model = "bge-large-en"
     end
 
     fab!(:topic) { Fabricate(:topic) }
diff --git a/spec/lib/modules/embeddings/semantic_search_spec.rb b/spec/lib/modules/embeddings/semantic_search_spec.rb
index 47c88c4f..ad18da71 100644
--- a/spec/lib/modules/embeddings/semantic_search_spec.rb
+++ b/spec/lib/modules/embeddings/semantic_search_spec.rb
@@ -24,7 +24,7 @@ RSpec.describe DiscourseAi::Embeddings::SemanticSearch do
     after { described_class.clear_cache_for(query) }
 
     def stub_candidate_ids(candidate_ids)
-      DiscourseAi::Embeddings::VectorRepresentations::AllMpnetBaseV2
+      DiscourseAi::Embeddings::VectorRepresentations::BgeLargeEn
         .any_instance
         .expects(:asymmetric_topics_similarity_search)
         .returns(candidate_ids)
diff --git a/spec/lib/modules/embeddings/semantic_topic_query_spec.rb b/spec/lib/modules/embeddings/semantic_topic_query_spec.rb
index 911bb4d8..369632e1 100644
--- a/spec/lib/modules/embeddings/semantic_topic_query_spec.rb
+++ b/spec/lib/modules/embeddings/semantic_topic_query_spec.rb
@@ -12,7 +12,7 @@ describe DiscourseAi::Embeddings::EntryPoint do
       fab!(:target) { Fabricate(:topic) }
 
       def stub_semantic_search_with(results)
-        DiscourseAi::Embeddings::VectorRepresentations::AllMpnetBaseV2
+        DiscourseAi::Embeddings::VectorRepresentations::BgeLargeEn
           .any_instance
           .expects(:symmetric_topics_similarity_search)
           .returns(results.concat([target.id]))
diff --git a/spec/support/embeddings_generation_stubs.rb b/spec/support/embeddings_generation_stubs.rb
index 4724b57a..5d688b94 100644
--- a/spec/support/embeddings_generation_stubs.rb
+++ b/spec/support/embeddings_generation_stubs.rb
@@ -3,6 +3,7 @@
 class EmbeddingsGenerationStubs
   class << self
     def discourse_service(model, string, embedding)
+      model = "bge-large-en-v1.5" if model == "bge-large-en"
       WebMock
         .stub_request(
           :post,