feat(cli): Implement QueryResult type.

Author: Zhe Yu

src/vectorcode/chunking.py

@@ -31,6 +31,9 @@ class Chunk:
     def __str__(self):
         return self.text
 
+    def __hash__(self) -> int:
+        return hash(f"VectorCodeChunk({self.start}:{self.end}@{self.text})")
+
     def export_dict(self):
         if self.start is not None:
             return {

src/vectorcode/subcommands/query/types.py

@@ -0,0 +1,73 @@
+import heapq
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Literal, Optional, Sequence, Union
+
+import numpy
+
+from vectorcode.chunking import Chunk
+
+
+@dataclass
+class QueryResult:
+    """
+    The container for one single query result.
+
+    args:
+    - path: path to the file
+    - content: `vectorcode.chunking.Chunk` object that stores the chunk
+    - query: query messages used for the search
+    - scores: similarity scores for the corresponding query.
+    """
+
+    path: str
+    chunk: Chunk
+    query: Sequence[str]
+    scores: Sequence[float]
+
+    @classmethod
+    def merge(cls, *results: "QueryResult") -> "QueryResult":
+        """
+        Given the results of a single chunk/document from different queries, merge them into a single `QueryResult` object.
+        """
+        for i in range(len(results) - 1):
+            if (i < len(results) - 1) and not results[i].is_same_doc(results[i + 1]):
+                raise ValueError(
+                    f"The inputs are not the same chunk: {results[i]}, {results[i + 1]}"
+                )
+
+        return QueryResult(
+            path=results[0].path,
+            chunk=results[0].chunk,
+            query=sum((tuple(i.query) for i in results), start=tuple()),
+            scores=sum((tuple(i.scores) for i in results), start=tuple()),
+        )
+
+    @staticmethod
+    def group(
+        *results: "QueryResult",
+        key: Union[Literal["path"], Literal["chunk"]] = "path",
+        top_k: Optional[int] = None,
+    ) -> dict[Chunk | str, list["QueryResult"]]:
+        assert key in {"path", "chunk"}
+        grouped_result: dict[Chunk | str, list["QueryResult"]] = defaultdict(list)
+
+        for res in results:
+            grouped_result[getattr(res, key)].append(res)
+
+        if top_k and top_k > 0:
+            for group in grouped_result.keys():
+                grouped_result[group] = heapq.nlargest(top_k, grouped_result[group])
+        return grouped_result
+
+    def mean_score(self):
+        return float(numpy.mean(self.scores))
+
+    def __lt__(self, other: "QueryResult"):
+        return self.mean_score() < other.mean_score()
+
+    def __gt__(self, 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

@@ -0,0 +1,82 @@
+import pytest
+from tree_sitter import Point
+
+from vectorcode.chunking import Chunk
+from vectorcode.subcommands.query.types import QueryResult
+
+
+def make_dummy_chunk():
+    return QueryResult(
+        path="dummy1.py",
+        chunk=Chunk(
+            text="hello", start=Point(row=1, column=0), end=Point(row=1, column=4)
+        ),
+        query=["hello"],
+        scores=[0.9],
+    )
+
+
+def test_QueryResult_merge():
+    res1, res2 = (make_dummy_chunk(), make_dummy_chunk())
+    res2.query = ["bye"]
+    res2.scores = [0.1]
+
+    merged = QueryResult.merge(res1, res2)
+    assert merged.path == res1.path
+    assert merged.chunk == res1.chunk
+    assert merged.mean_score() == 0.5
+    assert merged.query == ("hello", "bye")
+
+
+def test_QueryResult_merge_failed():
+    res1, res2 = (make_dummy_chunk(), make_dummy_chunk())
+    res2.path = "dummy2.py"
+    with pytest.raises(ValueError):
+        QueryResult.merge(res1, res2)
+
+
+def test_QueryResult_group_by_path():
+    res1, res2 = (make_dummy_chunk(), make_dummy_chunk())
+    res2.chunk = Chunk(
+        "hello", start=Point(row=2, column=0), end=Point(row=2, column=4)
+    )
+    res2.query = ["bye"]
+    res2.scores = [0.1]
+
+    grouped_dict = QueryResult.group(res1, res2)
+    assert len(grouped_dict.keys()) == 1
+    assert len(grouped_dict["dummy1.py"]) == 2
+
+
+def test_QueryResult_group_by_chunk():
+    res1, res2 = (make_dummy_chunk(), make_dummy_chunk())
+    res2.query = ["bye"]
+    res2.scores = [0.1]
+
+    grouped_dict = QueryResult.group(res1, res2, key="chunk")
+    assert len(grouped_dict.keys()) == 1
+    assert len(grouped_dict[res1.chunk]) == 2
+
+
+def test_QueryResult_group_top_k():
+    res1, res2 = (make_dummy_chunk(), make_dummy_chunk())
+    res2.chunk = Chunk(
+        "hello", start=Point(row=2, column=0), end=Point(row=2, column=4)
+    )
+    res2.query = ["bye"]
+    res2.scores = [0.1]
+
+    grouped_dict = QueryResult.group(res1, res2, top_k=1)
+    assert len(grouped_dict.keys()) == 1
+    assert len(grouped_dict["dummy1.py"]) == 1
+    assert grouped_dict["dummy1.py"][0].query[0] == "hello"
+
+
+def test_QueryResult_lt():
+    res1, res2 = (make_dummy_chunk(), make_dummy_chunk())
+    res2.chunk = Chunk(
+        "hello", start=Point(row=2, column=0), end=Point(row=2, column=4)
+    )
+    res2.query = ["bye"]
+    res2.scores = [0.1]
+    assert res2 < res1