Check function in VectorSearchIndex.

Author: WaVEV

django_mongodb_backend/indexes.py

@@ -1,19 +1,27 @@
 import itertools
 from collections import defaultdict
 
+from django.core.checks import Error
 from django.db import NotSupportedError
 from django.db.models import (
+    BooleanField,
+    CharField,
+    DateField,
+    DateTimeField,
     DecimalField,
     FloatField,
     Index,
+    IntegerField,
+    TextField,
+    UUIDField,
 )
 from django.db.models.lookups import BuiltinLookup
 from django.db.models.sql.query import Query
 from django.db.models.sql.where import AND, XOR, WhereNode
 from pymongo import ASCENDING, DESCENDING
 from pymongo.operations import IndexModel, SearchIndexModel
 
-from django_mongodb_backend.fields import ArrayField
+from django_mongodb_backend.fields import ArrayField, ObjectIdAutoField, ObjectIdField
 
 from .query_utils import process_rhs
 
@@ -151,19 +159,66 @@ class VectorSearchIndex(SearchIndex):
     def __init__(self, *expressions, similarities="cosine", **kwargs):
         super().__init__(*expressions, **kwargs)
         # validate the similarities types
-        if isinstance(similarities, str):
-            self._check_similarity_functions([similarities])
-        else:
-            self._check_similarity_functions(similarities)
         self.similarities = similarities
 
-    def _check_similarity_functions(self, similarities):
+    def check(self, model):
+        errors = []
+        error_id_prefix = "django_mongodb_backend.indexes.VectorSearchIndex"
+        similarities = (
+            self.similarity if isinstance(self.similarities, list) else [self.similarities]
+        )
         for func in similarities:
             if func not in self.ALLOWED_SIMILARITY_FUNCTIONS:
-                raise ValueError(
+                errors.append(
                     f"{func} isn't a valid similarity function, options "
-                    f"'are {','.join(self.ALLOWED_SIMILARITY_FUNCTIONS)}"
+                    f"'are {','.join(self.ALLOWED_SIMILARITY_FUNCTIONS)}",
+                    obj=self,
+                    id=f"{error_id_prefix}.E003",
+                )
+        for field_name, _ in self.fields_orders:
+            field_ = model._meta.get_field(field_name)
+            if isinstance(field_, ArrayField):
+                try:
+                    int(field_.size)
+                except (ValueError, TypeError):
+                    errors.append(
+                        Error(
+                            "Atlas vector search requires size.",
+                            obj=self,
+                            id=f"{error_id_prefix}.E001",
+                        )
+                    )
+                if not isinstance(field_.base_field, FloatField | DecimalField):
+                    errors.append(
+                        Error(
+                            "Base type must be Float or Decimal.",
+                            obj=self,
+                            id=f"{error_id_prefix}.E002",
+                        )
+                    )
+            # filter - for fields that contain boolean, date, objectId,
+            # numeric, string, or UUID values. Reference:
+            # https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-type/#atlas-vector-search-index-fields
+            elif not isinstance(
+                field_,
+                BooleanField
+                | IntegerField
+                | DateField
+                | DateTimeField
+                | CharField
+                | TextField
+                | UUIDField
+                | ObjectIdField
+                | ObjectIdAutoField,
+            ):
+                errors.append(
+                    Error(
+                        f"Unsupported filter of type {field_.get_internal_type()}.",
+                        obj=self,
+                        id="django_mongodb_backend.indexes.VectorSearchIndex.E003",
+                    )
                 )
+        return errors
 
     def deconstruct(self):
         path, args, kwargs = super().deconstruct()
@@ -198,16 +253,7 @@ def get_pymongo_index_model(
                     }
                 )
             else:
-                field_type = field_.db_type(schema_editor.connection)
-                search_type = self.search_index_data_types(field_, field_type)
-                # filter - for fields that contain boolean, date, objectId, numeric,
-                # string, or UUID values. Reference:
-                # https://www.mongodb.com/docs/atlas/atlas-vector-search/vector-search-type/#atlas-vector-search-index-fields
-                if search_type in ("number", "string", "boolean", "objectId", "uuid", "date"):
-                    mappings["type"] = "filter"
-                else:
-                    field_type = field_.get_internal_type()
-                    raise ValueError(f"Unsupported filter of type {field_type}.")
+                mappings["type"] = "filter"
             fields.append(mappings)
         return SearchIndexModel(definition={"fields": fields}, name=self.name, type="vectorSearch")