mongoid 9.0.12 → 9.1.0

data/spec/mongoid/search_indexable_spec.rb

@@ -28,7 +28,7 @@ class SearchIndexHelper
 
   # Wait until all of the indexes with the given names are absent from the
   # search index list.
-  def wait_for_absense_of(*names)
+  def wait_for_absence_of(*names)
     names.flatten.each do |name|
       timeboxed_wait do
         break if collection.search_indexes(name: name).empty?
@@ -61,25 +61,30 @@ class SearchIndexHelper
   end
 end
 
-# rubocop:disable RSpec/MultipleMemoizedHelpers
 describe Mongoid::SearchIndexable do
-  before do
-    skip "#{described_class} requires at Atlas environment (set ATLAS_URI)" if ENV['ATLAS_URI'].nil?
-  end
+  # Unit tests — no Atlas connection required.
 
-  let(:model) do
-    Class.new do
-      include Mongoid::Document
-      store_in collection: BSON::ObjectId.new.to_s
+  describe '.search_index_specs' do
+    let(:model) do
+      Class.new do
+        include Mongoid::Document
 
-      search_index mappings: { dynamic: false }
-      search_index :with_dynamic_mappings, mappings: { dynamic: true }
+        store_in collection: BSON::ObjectId.new.to_s
+        search_index mappings: { dynamic: false }
+        search_index :with_dynamic_mappings, mappings: { dynamic: true }
+      end
     end
-  end
 
-  let(:helper) { SearchIndexHelper.new(model) }
+    let(:vector_model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        vector_search_index fields: [ { type: 'vector', path: 'embedding', numDimensions: 1536, similarity: 'cosine' } ]
+        vector_search_index :named_vector_index, fields: [ { type: 'vector', path: 'other', numDimensions: 768, similarity: 'euclidean' } ]
+      end
+    end
 
-  describe '.search_index_specs' do
     context 'when no search indexes have been defined' do
       it 'has no search index specs' do
         expect(Person.search_index_specs).to be_empty
@@ -88,60 +93,460 @@ describe Mongoid::SearchIndexable do
 
     context 'when search indexes have been defined' do
       it 'has search index specs' do
-        expect(model.search_index_specs).to be == [
+        expect(model.search_index_specs).to eq [
           { definition: { mappings: { dynamic: false } } },
           { name: 'with_dynamic_mappings', definition: { mappings: { dynamic: true } } }
         ]
       end
     end
+
+    context 'when vector search indexes have been defined' do
+      it 'has search index specs with vectorSearch type' do
+        expect(vector_model.search_index_specs).to eq [
+          { type: 'vectorSearch', definition: { fields: [ { type: 'vector', path: 'embedding', numDimensions: 1536, similarity: 'cosine' } ] } },
+          { type: 'vectorSearch', name: 'named_vector_index', definition: { fields: [ { type: 'vector', path: 'other', numDimensions: 768, similarity: 'euclidean' } ] } }
+        ]
+      end
+    end
+  end
+
+  describe '.vector_search_index' do
+    let(:vector_model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        field :embedding, type: Array
+        vector_search_index fields: [ { type: 'vector', path: 'embedding', numDimensions: 3, similarity: 'cosine' } ]
+      end
+    end
+
+    it 'adds a vector_search_score field to the model' do
+      expect(vector_model.fields).to have_key('vector_search_score')
+    end
+
+    it 'marks vector_search_score as readonly' do
+      expect(vector_model.readonly_attributes).to include('vector_search_score')
+    end
+
+    it 'defines vector_search_score only once across multiple declarations' do
+      model = Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        vector_search_index :idx1, fields: [ { type: 'vector', path: 'a', numDimensions: 3, similarity: 'cosine' } ]
+        vector_search_index :idx2, fields: [ { type: 'vector', path: 'b', numDimensions: 3, similarity: 'cosine' } ]
+      end
+      expect(model.fields.count { |k, _| k == 'vector_search_score' }).to eq 1
+    end
+
+    context 'when only a regular search_index is declared' do
+      it 'does not add vector_search_score' do
+        model = Class.new do
+          include Mongoid::Document
+
+          store_in collection: BSON::ObjectId.new.to_s
+          search_index mappings: { dynamic: true }
+        end
+        expect(model.fields).not_to have_key('vector_search_score')
+      end
+    end
   end
 
-  context 'when needing to first create search indexes' do
-    let(:requested_definitions) { model.search_index_specs.map { |spec| spec[:definition].with_indifferent_access } }
-    let(:index_names) { model.create_search_indexes }
-    let(:actual_indexes) { helper.wait_for(*index_names) }
-    let(:actual_definitions) { actual_indexes.map { |i| i['latestDefinition'] } }
+  describe '.vector_search argument validation' do
+    let(:no_index_model) do
+      Class.new do
+        include Mongoid::Document
 
-    describe '.create_search_indexes' do
-      it 'creates the indexes' do
-        expect(actual_definitions).to be == requested_definitions
+        store_in collection: BSON::ObjectId.new.to_s
       end
     end
 
-    describe '.search_indexes' do
-      before { actual_indexes } # wait for the indices to be created
+    let(:multi_index_model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        vector_search_index :idx1, fields: [ { type: 'vector', path: 'a', numDimensions: 3, similarity: 'cosine' } ]
+        vector_search_index :idx2, fields: [ { type: 'vector', path: 'b', numDimensions: 3, similarity: 'cosine' } ]
+      end
+    end
 
-      let(:queried_definitions) { model.search_indexes.map { |i| i['latestDefinition'] } }
+    let(:single_index_model) do
+      Class.new do
+        include Mongoid::Document
 
-      it 'queries the available search indexes' do
-        expect(queried_definitions).to be == requested_definitions
+        store_in collection: BSON::ObjectId.new.to_s
+        field :embedding, type: Array
+        vector_search_index fields: [ { type: 'vector', path: 'embedding', numDimensions: 3, similarity: 'cosine' } ]
       end
     end
 
-    describe '.remove_search_index' do
-      let(:target_index) { actual_indexes.first }
+    it 'raises ArgumentError when no vector search indexes are declared' do
+      expect { no_index_model.vector_search([ 0.1, 0.2, 0.3 ]) }
+        .to raise_error(ArgumentError, /No vector search indexes declared/)
+    end
+
+    it 'raises ArgumentError when multiple indexes exist and none is specified' do
+      expect { multi_index_model.vector_search([ 0.1, 0.2, 0.3 ]) }
+        .to raise_error(ArgumentError, /multiple vector search indexes/)
+    end
+
+    it 'raises ArgumentError when the specified index name does not exist' do
+      expect { single_index_model.vector_search([ 0.1, 0.2, 0.3 ], index: 'nonexistent') }
+        .to raise_error(ArgumentError, /No vector search index 'nonexistent'/)
+    end
+  end
+
+  describe '#vector_search argument validation' do
+    let(:model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        field :embedding, type: Array
+        vector_search_index fields: [ { type: 'vector', path: 'embedding', numDimensions: 3, similarity: 'cosine' } ]
+      end
+    end
+
+    it 'raises ArgumentError when the vector field is nil' do
+      doc = model.new(embedding: nil)
+      expect { doc.vector_search }.to raise_error(ArgumentError, /embedding is nil/)
+    end
+  end
+
+  describe '#auto_embed_search argument validation' do
+    let(:model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        auto_embed_field :description, model: 'voyage-4'
+      end
+    end
+
+    it 'raises ArgumentError when the text field is nil' do
+      doc = model.new(description: nil)
+      expect { doc.auto_embed_search }.to raise_error(ArgumentError, /description is nil/)
+    end
+
+    it 'raises ArgumentError when no auto-embed indexes are declared on the model' do
+      bare_model = Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        field :description, type: String
+      end
+      doc = bare_model.new(description: 'hello')
+      expect { doc.auto_embed_search }.to raise_error(ArgumentError, /No auto-embed indexes declared/)
+    end
+
+    it 'raises ArgumentError when multiple indexes exist and none is specified' do
+      multi_model = Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        auto_embed_field :description, model: 'voyage-4', index: :idx1
+        auto_embed_field :summary,     model: 'voyage-4', index: :idx2
+      end
+      doc = multi_model.new(description: 'hello')
+      expect { doc.auto_embed_search }.to raise_error(ArgumentError, /multiple auto-embed indexes/)
+    end
+  end
+
+  describe '.auto_embed_search argument validation' do
+    let(:no_index_model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+      end
+    end
+
+    let(:single_embed_model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        auto_embed_field :description, model: 'voyage-4'
+      end
+    end
+
+    let(:multi_embed_model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        auto_embed_field :description, model: 'voyage-4', index: :idx1
+        auto_embed_field :summary,     model: 'voyage-4', index: :idx2
+      end
+    end
+
+    it 'raises ArgumentError when no auto-embed indexes are declared' do
+      expect { no_index_model.auto_embed_search('hello') }
+        .to raise_error(ArgumentError, /No auto-embed indexes declared/)
+    end
+
+    it 'raises ArgumentError when multiple indexes exist and none is specified' do
+      expect { multi_embed_model.auto_embed_search('hello') }
+        .to raise_error(ArgumentError, /multiple auto-embed indexes/)
+    end
+
+    it 'raises ArgumentError when the specified index name does not exist' do
+      expect { single_embed_model.auto_embed_search('hello', index: 'nonexistent') }
+        .to raise_error(ArgumentError, /No auto-embed index 'nonexistent'/)
+    end
+
+    it 'raises ArgumentError when a model with only vector (non-autoEmbed) indexes is used' do
+      model = Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        vector_search_index fields: [ { type: 'vector', path: 'emb', numDimensions: 3, similarity: 'cosine' } ]
+      end
+      expect { model.auto_embed_search('hello') }
+        .to raise_error(ArgumentError, /No auto-embed indexes declared/)
+    end
+  end
+
+  # Atlas integration tests — skipped when ATLAS_URI is not set.
+
+  context 'Atlas integration' do
+    before do
+      skip "#{described_class} requires an Atlas environment (set ATLAS_URI)" if ENV['ATLAS_URI'].nil?
+    end
+
+    let(:model) do
+      Class.new do
+        include Mongoid::Document
+
+        store_in collection: BSON::ObjectId.new.to_s
+        search_index mappings: { dynamic: false }
+        search_index :with_dynamic_mappings, mappings: { dynamic: true }
+      end
+    end
+
+    let(:helper) { SearchIndexHelper.new(model) }
+
+    context 'when needing to first create search indexes' do
+      let(:requested_definitions) { model.search_index_specs.map { |spec| spec[:definition].with_indifferent_access } }
+      let(:index_names) { model.create_search_indexes }
+      let(:actual_indexes) { helper.wait_for(*index_names) }
+      let(:actual_definitions) { actual_indexes.map { |i| i['latestDefinition'] } }
+
+      describe '.create_search_indexes' do
+        it 'creates the indexes' do
+          expect(actual_definitions).to eq requested_definitions
+        end
+      end
+
+      describe '.search_indexes' do
+        before { actual_indexes } # wait for the indices to be created
+
+        let(:queried_definitions) { model.search_indexes.map { |i| i['latestDefinition'] } }
+
+        it 'queries the available search indexes' do
+          expect(queried_definitions).to eq requested_definitions
+        end
+      end
+
+      describe '.remove_search_index' do
+        let(:target_index) { actual_indexes.first }
+
+        before do
+          model.remove_search_index id: target_index['id']
+          helper.wait_for_absence_of target_index['name']
+        end
+
+        it 'removes the requested index' do
+          expect(model.search_indexes(id: target_index['id'])).to be_empty
+        end
+      end
+
+      describe '.remove_search_indexes' do
+        before do
+          actual_indexes # wait for the indexes to be created
+          model.remove_search_indexes
+          helper.wait_for_absence_of(actual_indexes.map { |i| i['name'] })
+        end
+
+        it 'removes the indexes' do
+          expect(model.search_indexes).to be_empty
+        end
+      end
+    end
+
+    context 'with a vector search index' do
+      let(:vector_model) do
+        Class.new do
+          include Mongoid::Document
+
+          store_in collection: BSON::ObjectId.new.to_s
+          field :embedding, type: Array
+          vector_search_index fields: [ { type: 'vector', path: 'embedding', numDimensions: 3, similarity: 'cosine' } ]
+        end
+      end
+
+      let(:vector_helper) { SearchIndexHelper.new(vector_model) }
+
+      # Three orthogonal unit vectors as a minimal, predictable dataset.
+      let!(:doc_a) { vector_model.create!(embedding: [ 1.0, 0.0, 0.0 ]) }
+      let!(:doc_b) { vector_model.create!(embedding: [ 0.0, 1.0, 0.0 ]) }
+      let!(:doc_c) { vector_model.create!(embedding: [ 0.0, 0.0, 1.0 ]) }
 
       before do
-        model.remove_search_index id: target_index['id']
-        helper.wait_for_absense_of target_index['name']
+        names = vector_model.create_search_indexes
+        vector_helper.wait_for(*names)
+      end
+
+      describe '.vector_search' do
+        let(:results) { vector_model.vector_search([ 1.0, 0.0, 0.0 ], limit: 3) }
+
+        it 'returns Mongoid document instances' do
+          expect(results).to all(be_a(vector_model))
+        end
+
+        it 'populates vector_search_score on each result' do
+          expect(results.map(&:vector_search_score)).to all(be_a(Float))
+        end
+
+        it 'excludes the vector field from results' do
+          expect(results.map(&:embedding)).to all(be_nil)
+        end
+
+        it 'orders results by descending score' do
+          scores = results.map(&:vector_search_score)
+          expect(scores).to eq(scores.sort.reverse)
+        end
       end
 
-      it 'removes the requested index' do
-        expect(model.search_indexes(id: target_index['id'])).to be_empty
+      describe '#vector_search' do
+        let(:results) { doc_a.vector_search(limit: 3) }
+
+        it 'returns Mongoid document instances' do
+          expect(results).to all(be_a(vector_model))
+        end
+
+        it 'populates vector_search_score on each result' do
+          expect(results.map(&:vector_search_score)).to all(be_a(Float))
+        end
+
+        it 'excludes the source document from results' do
+          expect(results.map(&:id)).not_to include(doc_a.id)
+        end
       end
     end
 
-    describe '.remove_search_indexes' do
+    context 'with an auto-embed search index' do
       before do
-        actual_indexes # wait for the indexes to be created
-        model.remove_search_indexes
-        helper.wait_for_absense_of(actual_indexes.map { |i| i['name'] })
+        skip 'per DRIVERS-3315: "Integration tests can be deferred to follow up work to complete once a CI-compatible server is available with this feature (autoembedding indexes)"'
       end
 
-      it 'removes the indexes' do
-        expect(model.search_indexes).to be_empty
+      let(:embed_model) do
+        Class.new do
+          include Mongoid::Document
+
+          store_in collection: BSON::ObjectId.new.to_s
+          auto_embed_field :description, model: 'voyage-4'
+        end
+      end
+
+      let(:embed_helper) { SearchIndexHelper.new(embed_model) }
+
+      context 'when creating an auto-embed index' do
+        before { embed_helper } # ensure collection is set up before indexes are created
+
+        let(:index_names) { embed_model.create_search_indexes }
+        let(:actual_indexes) { embed_helper.wait_for(*index_names) }
+
+        describe '.create_search_indexes' do
+          it 'creates the auto-embed index' do
+            expect(actual_indexes).not_to be_empty
+          end
+
+          it 'creates an index with the autoEmbed field type' do
+            field_types = actual_indexes
+                          .flat_map { |i| i.dig('latestDefinition', 'fields') }
+                          .map { |f| f['type'] }
+            expect(field_types).to include('autoEmbed')
+          end
+        end
+
+        describe '.auto_embed_search' do
+          before do
+            actual_indexes # wait for index to be ready
+            embed_model.create(description: 'machine learning and neural networks')
+            embed_model.create(description: 'recipe for chocolate cake')
+            embed_model.create(description: 'deep learning for natural language processing')
+            # Atlas may need a moment to index new documents
+            sleep 5
+          end
+
+          it 'returns documents ranked by text similarity' do
+            results = embed_model.auto_embed_search('AI and machine learning', limit: 2)
+            expect(results).not_to be_empty
+            expect(results.first).to be_a(embed_model)
+            expect(results.first.vector_search_score).to be_a(Float)
+          end
+
+          it 'does not return more documents than the limit' do
+            results = embed_model.auto_embed_search('AI', limit: 1)
+            expect(results.size).to be <= 1
+          end
+
+          it 'supports exact nearest-neighbor search' do
+            results = embed_model.auto_embed_search('neural networks', exact: true, limit: 2)
+            expect(results).not_to be_empty
+          end
+        end
+
+        describe '#auto_embed_search' do
+          before do
+            actual_indexes # wait for index to be ready
+            embed_model.create(description: 'machine learning and neural networks')
+            embed_model.create(description: 'recipe for chocolate cake')
+            sleep 5
+          end
+
+          it 'returns similar documents excluding self' do
+            doc = embed_model.create(description: 'deep learning research')
+            sleep 5 # wait for the new document to be indexed
+            results = doc.auto_embed_search(limit: 5)
+            expect(results).not_to be_empty
+            expect(results.map(&:id)).not_to include(doc.id)
+          end
+        end
+      end
+
+      context 'with numDimensions and quantization options' do
+        let(:embed_model_with_opts) do
+          Class.new do
+            include Mongoid::Document
+
+            store_in collection: BSON::ObjectId.new.to_s
+            auto_embed_field :description,
+                             model: 'voyage-4',
+                             num_dimensions: 512,
+                             quantization: 'scalar'
+          end
+        end
+
+        let(:opts_helper) { SearchIndexHelper.new(embed_model_with_opts) }
+
+        before { opts_helper } # initialize (drop/create collection) before index creation
+
+        it 'creates the index with the specified numDimensions and quantization' do
+          index_names = embed_model_with_opts.create_search_indexes
+          actual_indexes = opts_helper.wait_for(*index_names)
+
+          field_spec = actual_indexes
+                       .flat_map { |i| i.dig('latestDefinition', 'fields') }
+                       .find { |f| f['type'] == 'autoEmbed' }
+
+          expect(field_spec['numDimensions']).to eq 512
+          expect(field_spec['quantization']).to eq 'scalar'
+        end
       end
     end
   end
 end
-# rubocop:enable RSpec/MultipleMemoizedHelpers