API refactor

venv/lib/python3.12/site-packages/redis/commands/search/aggregation.py

@@ -1,5 +1,7 @@
 from typing import List, Union
 
+from redis.commands.search.dialect import DEFAULT_DIALECT
+
 FIELDNAME = object()
 
 
@@ -24,7 +26,7 @@ class Reducer:
 
     NAME = None
 
-    def __init__(self, *args: List[str]) -> None:
+    def __init__(self, *args: str) -> None:
         self._args = args
         self._field = None
         self._alias = None
@@ -110,9 +112,11 @@ class AggregateRequest:
         self._with_schema = False
         self._verbatim = False
         self._cursor = []
-        self._dialect = None
+        self._dialect = DEFAULT_DIALECT
+        self._add_scores = False
+        self._scorer = "TFIDF"
 
-    def load(self, *fields: List[str]) -> "AggregateRequest":
+    def load(self, *fields: str) -> "AggregateRequest":
         """
         Indicate the fields to be returned in the response. These fields are
         returned in addition to any others implicitly specified.
@@ -219,7 +223,7 @@ class AggregateRequest:
         self._aggregateplan.extend(_limit.build_args())
         return self
 
-    def sort_by(self, *fields: List[str], **kwargs) -> "AggregateRequest":
+    def sort_by(self, *fields: str, **kwargs) -> "AggregateRequest":
         """
         Indicate how the results should be sorted. This can also be used for
         *top-N* style queries
@@ -292,6 +296,24 @@ class AggregateRequest:
         self._with_schema = True
         return self
 
+    def add_scores(self) -> "AggregateRequest":
+        """
+        If set, includes the score as an ordinary field of the row.
+        """
+        self._add_scores = True
+        return self
+
+    def scorer(self, scorer: str) -> "AggregateRequest":
+        """
+        Use a different scoring function to evaluate document relevance.
+        Default is `TFIDF`.
+
+        :param scorer: The scoring function to use
+                       (e.g. `TFIDF.DOCNORM` or `BM25`)
+        """
+        self._scorer = scorer
+        return self
+
     def verbatim(self) -> "AggregateRequest":
         self._verbatim = True
         return self
@@ -315,12 +337,19 @@ class AggregateRequest:
         if self._verbatim:
             ret.append("VERBATIM")
 
+        if self._scorer:
+            ret.extend(["SCORER", self._scorer])
+
+        if self._add_scores:
+            ret.append("ADDSCORES")
+
         if self._cursor:
             ret += self._cursor
 
         if self._loadall:
             ret.append("LOAD")
             ret.append("*")
+
         elif self._loadfields:
             ret.append("LOAD")
             ret.append(str(len(self._loadfields)))