make similarities required

Author: timgraham

django_mongodb_backend/indexes.py

@@ -167,7 +167,7 @@ class VectorSearchIndex(SearchIndex):
     VALID_FIELD_TYPES = frozenset(("boolean", "date", "number", "objectId", "string", "uuid"))
     _error_id_prefix = "django_mongodb_backend.indexes.VectorSearchIndex"
 
-    def __init__(self, *, fields=(), similarities="cosine", name=None):
+    def __init__(self, *, fields=(), name=None, similarities):
         super().__init__(fields=fields, name=name)
         self.similarities = similarities
         self._multiple_similarities = isinstance(similarities, tuple | list)
@@ -248,8 +248,7 @@ def check(self, model, connection):
 
     def deconstruct(self):
         path, args, kwargs = super().deconstruct()
-        if self.similarities != "cosine":
-            kwargs["similarities"] = self.similarities
+        kwargs["similarities"] = self.similarities
         return path, args, kwargs
 
     def get_pymongo_index_model(

docs/source/ref/models/indexes.rst

@@ -31,24 +31,24 @@ model has multiple indexes).
 ``VectorSearchIndex``
 =====================
 
-.. class:: VectorSearchIndex(fields=(), similarities="cosine", name=None)
+.. class:: VectorSearchIndex(*, fields=(), name=None, similarities)
 
 .. versionadded:: 5.2.0b0
 
 A subclass of :class:`SearchIndex` that creates a :doc:`vector search index
 <atlas:atlas-vector-search/vector-search-type>` on the given field(s).
 
-The index should reference at least one vector field: an :class:`.ArrayField`
+The index must reference at least one vector field: an :class:`.ArrayField`
 with a :attr:`~.ArrayField.base_field` of :class:`~django.db.models.FloatField`
-or :class:`~django.db.models.IntegerField`. It cannot reference an
-:class:`.ArrayField` of any other type. Each :class:`.ArrayField` must have a
-:attr:`~.ArrayField.size`.
+or :class:`~django.db.models.IntegerField` and a :attr:`~.ArrayField.size`. It
+cannot reference an :class:`.ArrayField` of any other type.
 
 It may also have other fields to filter on, provided the field stores
 ``boolean``, ``date``, ``objectId``, ``numeric``, ``string``, or ``uuid``.
 
-Available values for ``similarities`` are ``"cosine"``, ``"dotProduct"``, and
-``"euclidean"`` (see :ref:`atlas:avs-similarity-functions`). You can provide
-this value either a string, in which case that value will be applied to all
-vector fields, or a list or tuple of values with a similarity corresponding to
-each vector field.
+Available values for the required ``similarities`` keyword argument are
+``"cosine"``, ``"dotProduct"``, and ``"euclidean"`` (see
+:ref:`atlas:avs-similarity-functions` for how to choose). You can provide this
+value either a string, in which case that value will be applied to all vector
+fields, or a list or tuple of values with a similarity corresponding to each
+vector field.

tests/indexes_/test_checks.py

@@ -44,7 +44,7 @@ class Article(models.Model):
             vector = ArrayField(models.FloatField(), size=10)
 
             class Meta:
-                indexes = [VectorSearchIndex(fields=["title", "vector"])]
+                indexes = [VectorSearchIndex(fields=["title", "vector"], similarities="cosine")]
 
         errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
         self.assertEqual(
@@ -72,7 +72,7 @@ class Article(models.Model):
             title_embedded = ArrayField(models.FloatField())
 
             class Meta:
-                indexes = [VectorSearchIndex(fields=["title_embedded"])]
+                indexes = [VectorSearchIndex(fields=["title_embedded"], similarities="cosine")]
 
         errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
         self.assertEqual(
@@ -91,7 +91,7 @@ class Article(models.Model):
             title_embedded = ArrayField(models.CharField(), size=30)
 
             class Meta:
-                indexes = [VectorSearchIndex(fields=["title_embedded"])]
+                indexes = [VectorSearchIndex(fields=["title_embedded"], similarities="cosine")]
 
         errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
         self.assertEqual(
@@ -112,7 +112,7 @@ class Article(models.Model):
             vector = ArrayField(models.FloatField(), size=10)
 
             class Meta:
-                indexes = [VectorSearchIndex(fields=["data", "vector"])]
+                indexes = [VectorSearchIndex(fields=["data", "vector"], similarities="cosine")]
 
         errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
         self.assertEqual(
@@ -127,7 +127,7 @@ class Meta:
             ],
         )
 
-    def test_invalid_number_similarity_function_singular(self):
+    def test_fields_and_similarities_mismatch(self):
         class Article(models.Model):
             vector = ArrayField(models.FloatField(), size=10)
 
@@ -153,44 +153,17 @@ class Meta:
             ],
         )
 
-    def test_invalid_number_similarity_function_plural(self):
-        class Article(models.Model):
-            vector1 = ArrayField(models.FloatField(), size=10)
-            vector2 = ArrayField(models.FloatField(), size=10)
-
-            class Meta:
-                indexes = [
-                    VectorSearchIndex(
-                        fields=["vector1", "vector2"],
-                        similarities=["dotProduct"],
-                    )
-                ]
-
-        errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
-        self.assertEqual(
-            errors,
-            [
-                checks.Error(
-                    "VectorSearchIndex requires the same number of similarities "
-                    "and vector fields; Article has 2 ArrayField(s) but similarities "
-                    "has 1 element(s).",
-                    id="django_mongodb_backend.indexes.VectorSearchIndex.E005",
-                    obj=Article,
-                ),
-            ],
-        )
-
     def test_simple(self):
         class Article(models.Model):
             vector = ArrayField(models.FloatField(), size=10)
 
             class Meta:
-                indexes = [VectorSearchIndex(fields=["vector"])]
+                indexes = [VectorSearchIndex(fields=["vector"], similarities="cosine")]
 
         errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
         self.assertEqual(errors, [])
 
-    def test_all_valid_fields(self):
+    def test_valid_fields(self):
         class Data(EmbeddedModel):
             integer = models.IntegerField()
 
@@ -216,6 +189,7 @@ class Meta:
                             "boolean",
                             "date",
                         ],
+                        similarities="cosine",
                     )
                 ]
 
@@ -227,7 +201,11 @@ class NoSearchVectorModel(models.Model):
             text = models.CharField(max_length=100)
 
             class Meta:
-                indexes = [VectorSearchIndex(name="recent_test_idx", fields=["text"])]
+                indexes = [
+                    VectorSearchIndex(
+                        name="recent_test_idx", fields=["text"], similarities="cosine"
+                    )
+                ]
 
         errors = checks.run_checks(app_configs=self.apps.get_app_configs(), databases={"default"})
         self.assertEqual(

tests/indexes_/test_search_indexes.py

@@ -26,7 +26,7 @@ def test_search_index_not_created(self):
             editor.remove_index(index=index, model=SearchIndexTestModel)
 
     def test_vector_index_not_created(self):
-        index = VectorSearchIndex(name="recent_test_idx", fields=["number"])
+        index = VectorSearchIndex(name="recent_test_idx", fields=["number"], similarities="cosine")
         with connection.schema_editor() as editor, self.assertNumQueries(0):
             editor.add_index(index=index, model=SearchIndexTestModel)
         self.assertNotIn(
@@ -67,27 +67,18 @@ def test_no_extra_kargs(self):
         with self.assertRaisesMessage(TypeError, msg):
             VectorSearchIndex(condition="")
 
+    def test_no_similarities(self):
+        msg = (
+            "VectorSearchIndex.__init__() missing 1 required keyword-only argument: 'similarities'"
+        )
+        with self.assertRaisesMessage(TypeError, msg):
+            VectorSearchIndex(name="recent_test_idx", fields=["number"])
+
     def test_deconstruct(self):
-        index = VectorSearchIndex(name="recent_test_idx", fields=["number"])
+        index = VectorSearchIndex(name="recent_test_idx", fields=["number"], similarities="cosine")
         name, args, kwargs = index.deconstruct()
-        self.assertEqual(kwargs, {"name": "recent_test_idx", "fields": ["number"]})
-        new = VectorSearchIndex(*args, **kwargs)
-        self.assertEqual(new.similarities, index.similarities)
-
-    def test_deconstruct_with_similarities(self):
-        index = VectorSearchIndex(
-            name="recent_test_idx",
-            fields=["number", "char"],
-            similarities=["cosine", "dotProduct"],
-        )
-        path, args, kwargs = index.deconstruct()
         self.assertEqual(
-            kwargs,
-            {
-                "name": "recent_test_idx",
-                "fields": ["number", "char"],
-                "similarities": ["cosine", "dotProduct"],
-            },
+            kwargs, {"name": "recent_test_idx", "fields": ["number"], "similarities": "cosine"}
         )
         new = VectorSearchIndex(*args, **kwargs)
         self.assertEqual(new.similarities, index.similarities)
@@ -195,7 +186,7 @@ def test_valid_fields(self):
 @skipUnlessDBFeature("supports_atlas_search")
 class VectorSearchIndexSchemaTests(SchemaAssertionMixin, TestCase):
     def test_simple(self):
-        index = VectorSearchIndex(name="recent_test_idx", fields=["integer"])
+        index = VectorSearchIndex(name="recent_test_idx", fields=["integer"], similarities="cosine")
         with connection.schema_editor() as editor:
             self.assertAddRemoveIndex(editor, index=index, model=SearchIndexTestModel)
 
@@ -212,6 +203,7 @@ def test_multiple_fields(self):
                 "vector_float",
                 "vector_integer",
             ],
+            similarities="cosine",
         )
         with connection.schema_editor() as editor:
             editor.add_index(index=index, model=SearchIndexTestModel)