Allow OpenSearch embedding field name to be configurable

In order to support multiple embedding providers (for evaluating
different ones), we'll need to store these different embeddings in
different fields in OpenSearch.

So we'll need to be able to specify the embedding provider we're using
and use that to determine the field name.

Currently we're just supporting OpenAI, with a view to extending this
soon.

Author: jackbot

lib/search/chunked_content_repository.rb

@@ -141,14 +141,21 @@ def id_digest_hash(base_path, batch_size: 100)
       items
     end
 
-    def search_by_embedding(embedding, max_chunks:)
+    def search_by_embedding(embedding, max_chunks:, llm_provider:)
+      field_name = case llm_provider.to_sym
+                   when :openai
+                     :openai_embedding
+                   else
+                     raise "Unknown provider: #{llm_provider}"
+                   end
+
       response = client.search(
         index:,
         body: {
           size: max_chunks,
           query: {
             knn: {
-              openai_embedding: {
+              "#{field_name}": {
                 vector: embedding,
                 k: max_chunks,
               },

lib/search/results_for_question.rb

@@ -15,7 +15,11 @@ def self.call(question_message)
       metrics[:embedding_duration] = Clock.monotonic_time - embedding_start_time
 
       search_start_time = Clock.monotonic_time
-      results = ChunkedContentRepository.new.search_by_embedding(embedding, max_chunks:)
+      results = ChunkedContentRepository.new.search_by_embedding(
+        embedding,
+        max_chunks:,
+        llm_provider: provider,
+      )
       metrics[:search_duration] = Clock.monotonic_time - search_start_time
       metrics[:embedding_provider] = provider
 

spec/lib/search/chunked_content_repository_spec.rb

@@ -220,7 +220,11 @@
     end
 
     it "returns an array of Result objects" do
-      result = repository.search_by_embedding(openai_embedding, max_chunks: 10)
+      result = repository.search_by_embedding(
+        openai_embedding,
+        max_chunks: 10,
+        llm_provider: :openai,
+      )
       expected_attributes = chunked_content_records.first
                                                    .except(:openai_embedding)
                                                    .merge(score: a_value_between(0.9, 1))
@@ -229,12 +233,26 @@
       expect(result.first).to have_attributes(**expected_attributes)
     end
 
-    context "when there are more then the maxiumum chunks" do
+    it "raises an error if the llm provider is not recognised" do
+      expect {
+        repository.search_by_embedding(
+          openai_embedding,
+          max_chunks: 10,
+          llm_provider: :unknown,
+        )
+      }.to raise_error("Unknown provider: unknown")
+    end
+
+    context "when there are more than the maxiumum chunks" do
       let(:max_chunks) { 10 }
       let(:chunked_content_records) { build_list(:chunked_content_record, 11, openai_embedding:) }
 
       it "only returns the first max_chunks" do
-        result = repository.search_by_embedding(openai_embedding, max_chunks:)
+        result = repository.search_by_embedding(
+          openai_embedding,
+          max_chunks:,
+          llm_provider: :openai,
+        )
         expect(result.count).to eq max_chunks
       end
     end

spec/requests/admin/search_spec.rb

@@ -67,7 +67,7 @@
         it "includes score calculation and back links in links to results" do
           get admin_search_path, params: { search_text: }
 
-          results = Search::ChunkedContentRepository.new.search_by_embedding(openai_embedding, max_chunks: 2)
+          results = Search::ChunkedContentRepository.new.search_by_embedding(openai_embedding, max_chunks: 2, llm_provider: :openai)
           result = results.detect { |r| r.digest == chunk_to_find[:digest] }
           document_type_weight = Search::ResultsForQuestion::Reranker::DOCUMENT_TYPE_WEIGHTINGS[chunk_to_find[:document_type]]
           weighted_score = result.score * document_type_weight