Add Autocomplete operator.

Author: WaVEV

django_mongodb_backend/expressions/search.py

@@ -1,11 +1,115 @@
-from django.db.models import Expression
+from django.db.models import Expression, FloatField
+from django.db.models.expressions import F, Value
+
+
+def cast_as_value(value):
+    if value is None:
+        return None
+    return Value(value) if not hasattr(value, "resolve_expression") else value
+
+
+def cast_as_field(path):
+    return F(path) if isinstance(path, str) else path
 
 
 class SearchExpression(Expression):
-    """Base expression node for MongoDB Atlas `$search` stages."""
+    """Base expression node for MongoDB Atlas `$search` stages.
+
+    This class bridges Django's `Expression` API with the MongoDB Atlas
+    Search engine.  Subclasses produce the operator document placed under
+    **$search** and expose the stage to queryset methods such as
+    `annotate()`, `filter()`, or `order_by()`.
+    """
+
+    output_field = FloatField()
+
+    def __str__(self):
+        cls = self.identity[0]
+        kwargs = dict(self.identity[1:])
+        arg_str = ", ".join(f"{k}={v!r}" for k, v in kwargs.items())
+        return f"<{cls.__name__}({arg_str})>"
+
+    def __repr__(self):
+        return str(self)
+
+    def as_sql(self, compiler, connection):
+        return "", []
+
+    def _get_indexed_fields(self, mappings):
+        if isinstance(mappings, list):
+            for definition in mappings:
+                yield from self._get_indexed_fields(definition)
+        else:
+            for field, definition in mappings.get("fields", {}).items():
+                yield field
+                for path in self._get_indexed_fields(definition):
+                    yield f"{field}.{path}"
+
+    def _get_query_index(self, fields, compiler):
+        fields = set(fields)
+        for search_indexes in compiler.collection.list_search_indexes():
+            mappings = search_indexes["latestDefinition"]["mappings"]
+            indexed_fields = set(self._get_indexed_fields(mappings))
+            if mappings["dynamic"] or fields.issubset(indexed_fields):
+                return search_indexes["name"]
+        return "default"
 
+    def search_operator(self, compiler, connection):
+        raise NotImplementedError
 
-class SearchVector(SearchExpression):
+    def as_mql(self, compiler, connection):
+        index = self._get_query_index(self.get_search_fields(compiler, connection), compiler)
+        return {"$search": {**self.search_operator(compiler, connection), "index": index}}
+
+
+class SearchAutocomplete(SearchExpression):
     """
-    Atlas Search expression that performs vector similarity search on embedded vectors.
+    Atlas Search expression that matches input using the `autocomplete` operator.
+
+    This expression enables autocomplete behavior by querying against a field
+    indexed as `"type": "autocomplete"` in MongoDB Atlas. It can be used in
+    `filter()`, `annotate()` or any context that accepts a Django expression.
+
+    Example:
+        SearchAutocomplete("title", "harry", fuzzy={"maxEdits": 1})
+
+    Args:
+        path: The document path to search (as string or expression).
+        query: The input string to autocomplete.
+        fuzzy: Optional dictionary of fuzzy matching parameters.
+        token_order: Optional value for `"tokenOrder"`; controls sequential vs.
+            any-order token matching.
+        score: Optional expression to adjust score relevance (e.g., `{"boost": {"value": 5}}`).
+
+    Reference: https://www.mongodb.com/docs/atlas/atlas-search/autocomplete/
     """
+
+    def __init__(self, path, query, fuzzy=None, token_order=None, score=None):
+        self.path = cast_as_field(path)
+        self.query = cast_as_value(query)
+        self.fuzzy = cast_as_value(fuzzy)
+        self.token_order = cast_as_value(token_order)
+        self.score = score
+        super().__init__()
+
+    def get_source_expressions(self):
+        return [self.path, self.query, self.fuzzy, self.token_order]
+
+    def set_source_expressions(self, exprs):
+        self.path, self.query, self.fuzzy, self.token_order = exprs
+
+    def get_search_fields(self, compiler, connection):
+        return {self.path.as_mql(compiler, connection, as_path=True)}
+
+    def search_operator(self, compiler, connection):
+        params = {
+            "path": self.path.as_mql(compiler, connection, as_path=True),
+            "query": self.query.value,
+        }
+        if self.score is not None:
+            params["score"] = self.score.as_mql(compiler, connection)
+        if self.fuzzy is not None:
+            params["fuzzy"] = self.fuzzy.value
+        if self.token_order is not None:
+            params["tokenOrder"] = self.token_order.value
+        return {"autocomplete": params}