rename parameter

Author: Zhe Yu

src/vectorcode/subcommands/query/reranker/base.py

@@ -62,7 +62,7 @@ async def rerank(self, results: list[QueryResult]) -> list[str]:
         group_by = "path"
         if QueryInclude.chunk in self.configs.include:
             group_by = "chunk"
-        grouped_results = QueryResult.group(*results, key=group_by, top_k="auto")
+        grouped_results = QueryResult.group(*results, by=group_by, top_k="auto")
 
         scores: dict[Chunk | str, float] = {}
         for key in grouped_results.keys():

src/vectorcode/subcommands/query/types.py

@@ -46,14 +46,25 @@ def merge(cls, *results: "QueryResult") -> "QueryResult":
     @staticmethod
     def group(
         *results: "QueryResult",
-        key: Union[Literal["path"], Literal["chunk"]] = "path",
+        by: Union[Literal["path"], Literal["chunk"]] = "path",
         top_k: int | Literal["auto"] | None = None,
     ) -> dict[Chunk | str, list["QueryResult"]]:
-        assert key in {"path", "chunk"}
+        """
+        Group the query results based on `key`.
+
+        args:
+        - `by`: either "path" or "chunk"
+        - `top_k`: if set, only return the top k results for each group based on mean scores. If "auto", top k is decided by the mean number of results per group.
+
+        returns:
+        - a dictionary that maps either path or chunk to a list of `QueryResult` object.
+
+        """
+        assert by in {"path", "chunk"}
         grouped_result: dict[Chunk | str, list["QueryResult"]] = defaultdict(list)
 
         for res in results:
-            grouped_result[getattr(res, key)].append(res)
+            grouped_result[getattr(res, by)].append(res)
 
         if top_k == "auto":
             top_k = int(numpy.mean(tuple(len(i) for i in grouped_result.values())))
@@ -67,10 +78,16 @@ def mean_score(self):
         return float(numpy.mean(self.scores))
 
     def __lt__(self, other: "QueryResult"):
+        assert isinstance(other, QueryResult)
         return self.mean_score() < other.mean_score()
 
     def __gt__(self, other: "QueryResult"):
+        assert isinstance(other, QueryResult)
         return self.mean_score() > other.mean_score()
 
+    def __eq__(self, other: object, /) -> bool:
+        assert isinstance(other, QueryResult)
+        return self.mean_score() == other.mean_score()
+
     def is_same_doc(self, other: "QueryResult") -> bool:
         return self.path == other.path and self.chunk == other.chunk

tests/subcommands/query/test_types.py

@@ -53,7 +53,7 @@ def test_QueryResult_group_by_chunk():
     res2.query = ["bye"]
     res2.scores = [0.1]
 
-    grouped_dict = QueryResult.group(res1, res2, key="chunk")
+    grouped_dict = QueryResult.group(res1, res2, by="chunk")
     assert len(grouped_dict.keys()) == 1
     assert len(grouped_dict[res1.chunk]) == 2