refactor(cli): Use in-house QueryResult in the query subcommand and the rerankers.

Author: Zhe Yu

src/vectorcode/chunking.py

@@ -25,8 +25,8 @@ class Chunk:
     """
 
     text: str
-    start: Point
-    end: Point
+    start: Point | None = None
+    end: Point | None = None
 
     def __str__(self):
         return self.text
@@ -35,14 +35,20 @@ def __hash__(self) -> int:
         return hash(f"VectorCodeChunk({self.start}:{self.end}@{self.text})")
 
     def export_dict(self):
+        d: dict[str, str | dict[str, int]] = {"text": self.text}
         if self.start is not None:
-            return {
-                "text": self.text,
-                "start": {"row": self.start.row, "column": self.start.column},
-                "end": {"row": self.end.row, "column": self.end.column},
-            }
-        else:
-            return {"text": self.text}
+            d.update(
+                {
+                    "start": {"row": self.start.row, "column": self.start.column},
+                }
+            )
+        if self.end is not None:
+            d.update(
+                {
+                    "end": {"row": self.end.row, "column": self.end.column},
+                }
+            )
+        return d
 
 
 @dataclass

src/vectorcode/subcommands/query/__init__.py

@@ -5,10 +5,11 @@
 
 from chromadb import GetResult, Where
 from chromadb.api.models.AsyncCollection import AsyncCollection
-from chromadb.api.types import IncludeEnum
+from chromadb.api.types import IncludeEnum, QueryResult
 from chromadb.errors import InvalidCollectionException, InvalidDimensionException
+from tree_sitter import Point
 
-from vectorcode.chunking import StringChunker
+from vectorcode.chunking import Chunk, StringChunker
 from vectorcode.cli_utils import (
     Config,
     QueryInclude,
@@ -22,6 +23,7 @@
     get_embedding_function,
     verify_ef,
 )
+from vectorcode.subcommands.query import types as vectorcode_types
 from vectorcode.subcommands.query.reranker import (
     RerankerError,
     get_reranker,
@@ -30,14 +32,45 @@
 logger = logging.getLogger(name=__name__)
 
 
+def conver_query_results(
+    chroma_result: QueryResult, queries: list[str]
+) -> list[vectorcode_types.QueryResult]:
+    """Convert chromadb query result to in-house query results"""
+    assert chroma_result["documents"] is not None
+    assert chroma_result["distances"] is not None
+    assert chroma_result["metadatas"] is not None
+
+    chroma_results_list: list[vectorcode_types.QueryResult] = []
+    for q_i in range(len(queries)):
+        q = queries[q_i]
+        documents = chroma_result["documents"][q_i]
+        distances = chroma_result["distances"][q_i]
+        metadatas = chroma_result["metadatas"][q_i]
+        for doc, dist, meta in zip(documents, distances, metadatas):
+            chunk = Chunk(text=doc)
+            if meta["start"]:
+                chunk.start = Point(int(meta.get("start", 0)), 0)
+            if meta["end"]:
+                chunk.end = Point(int(meta.get("end", 0)) + 1, 0)
+            chroma_results_list.append(
+                vectorcode_types.QueryResult(
+                    chunk=chunk,
+                    path=str(meta.get("path", "")),
+                    query=(q,),
+                    scores=(-dist,),
+                )
+            )
+    return chroma_results_list
+
+
 async def get_query_result_files(
     collection: AsyncCollection, configs: Config
 ) -> list[str]:
     query_chunks = []
-    if configs.query:
-        chunker = StringChunker(configs)
-        for q in configs.query:
-            query_chunks.extend(str(i) for i in chunker.chunk(q))
+    assert configs.query, "Query messages cannot be empty."
+    chunker = StringChunker(configs)
+    for q in configs.query:
+        query_chunks.extend(str(i) for i in chunker.chunk(q))
 
     configs.query_exclude = [
         expand_path(i, True)
@@ -70,7 +103,7 @@ async def get_query_result_files(
         query_embeddings = get_embedding_function(configs)(query_chunks)
         if isinstance(configs.embedding_dims, int) and configs.embedding_dims > 0:
             query_embeddings = [e[: configs.embedding_dims] for e in query_embeddings]
-        results = await collection.query(
+        chroma_query_results: QueryResult = await collection.query(
             query_embeddings=query_embeddings,
             n_results=num_query,
             include=[
@@ -85,7 +118,9 @@ async def get_query_result_files(
         return []
 
     reranker = get_reranker(configs)
-    return await reranker.rerank(results)
+    return await reranker.rerank(
+        conver_query_results(chroma_query_results, configs.query)
+    )
 
 
 async def build_query_results(

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

@@ -1,13 +1,13 @@
 import heapq
 import logging
 from abc import ABC, abstractmethod
-from collections import defaultdict
-from typing import Any, DefaultDict, Optional, Sequence, cast
+from typing import Any
 
 import numpy
-from chromadb.api.types import QueryResult
 
+from vectorcode.chunking import Chunk
 from vectorcode.cli_utils import Config, QueryInclude
+from vectorcode.subcommands.query.types import QueryResult
 
 logger = logging.getLogger(name=__name__)
 
@@ -29,7 +29,7 @@ def __init__(self, configs: Config, **kwargs: Any):
             "'configs' should contain the query messages."
         )
         self.n_result = configs.n_result
-        self._raw_results: Optional[QueryResult] = None
+        self._raw_results: list[QueryResult] = []
 
     @classmethod
     def create(cls, configs: Config, **kwargs: Any):
@@ -46,53 +46,31 @@ def create(cls, configs: Config, **kwargs: Any):
             raise
 
     @abstractmethod
-    async def compute_similarity(
-        self, results: list[str], query_message: str
-    ) -> Sequence[float]:  # pragma: nocover
-        """Given a list of n results and 1 query message,
-        return a list-like object of length n that contains the similarity scores between
-        each item in `results` and the `query_message`.
-
-        A high similarity score means the strings are semantically similar to each other.
-        `query_message` will be loaded in the same order as they appear in `self.configs.query`.
-
-        If you need the raw query results from chromadb,
-        it'll be saved in `self._raw_results` before this method is called.
+    async def compute_similarity(self, results: list[QueryResult]):  # pragma: nocover
+        """
+        Modify the `QueryResult.scores` field IN-PLACE so that they contain the correct scores.
         """
         raise NotImplementedError
 
-    async def rerank(self, results: QueryResult | dict) -> list[str]:
-        if len(results["ids"]) == 0 or all(len(i) == 0 for i in results["ids"]):
+    async def rerank(self, results: list[QueryResult]) -> list[str]:
+        if len(results) == 0:
             return []
+        results = await self.compute_similarity(results)
 
-        self._raw_results = cast(QueryResult, results)
-        query_chunks = self.configs.query
-        assert query_chunks
-        assert results["metadatas"] is not None
-        assert results["documents"] is not None
-        documents: DefaultDict[str, list[float]] = defaultdict(list)
-        for query_chunk_idx in range(len(query_chunks)):
-            chunk_ids = results["ids"][query_chunk_idx]
-            chunk_metas = results["metadatas"][query_chunk_idx]
-            chunk_docs = results["documents"][query_chunk_idx]
-            scores = await self.compute_similarity(
-                chunk_docs, query_chunks[query_chunk_idx]
-            )
-            for i, score in enumerate(scores):
-                if QueryInclude.chunk in self.configs.include:
-                    documents[chunk_ids[i]].append(float(score))
-                else:
-                    documents[str(chunk_metas[i]["path"])].append(float(score))
+        group_by = "path"
+        if QueryInclude.chunk in self.configs.include:
+            group_by = "chunk"
+        grouped_results = QueryResult.group(*results, key=group_by, top_k="auto")
 
-        logger.debug("Document scores: %s", documents)
-        top_k = int(numpy.mean(tuple(len(i) for i in documents.values())))
-        for key in documents.keys():
-            documents[key] = heapq.nlargest(top_k, documents[key])
-
-        self._raw_results = None
+        scores: dict[Chunk | str, float] = {}
+        for key in grouped_results.keys():
+            scores[key] = float(
+                numpy.mean(tuple(i.mean_score() for i in grouped_results[key]))
+            )
 
-        return heapq.nlargest(
-            self.n_result,
-            documents.keys(),
-            key=lambda x: float(numpy.mean(documents[x])),
+        return list(
+            str(i)
+            for i in heapq.nlargest(
+                self.configs.n_result, grouped_results.keys(), key=lambda x: scores[x]
+            )
         )

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

@@ -1,8 +1,8 @@
-import asyncio
 import logging
 from typing import Any
 
 from vectorcode.cli_utils import Config
+from vectorcode.subcommands.query.types import QueryResult
 
 from .base import RerankerBase
 
@@ -34,8 +34,8 @@ def __init__(
         model_name = configs.reranker_params.pop("model_name_or_path")
         self.model = CrossEncoder(model_name, **configs.reranker_params)
 
-    async def compute_similarity(self, results: list[str], query_message: str):
-        scores = await asyncio.to_thread(
-            self.model.predict, [(chunk, query_message) for chunk in results]
-        )
-        return list(float(i) for i in scores)
+    async def compute_similarity(self, results: list[QueryResult]):
+        scores = self.model.predict([(str(res.chunk), res.query[0]) for res in results])
+
+        for res, score in zip(results, scores):
+            res.scores = (score,)

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

@@ -1,7 +1,8 @@
 import logging
-from typing import Any, Sequence
+from typing import Any
 
 from vectorcode.cli_utils import Config
+from vectorcode.subcommands.query.types import QueryResult
 
 from .base import RerankerBase
 
@@ -17,15 +18,8 @@ class NaiveReranker(RerankerBase):
     def __init__(self, configs: Config, **kwargs: Any):
         super().__init__(configs)
 
-    async def compute_similarity(
-        self, results: list[str], query_message: str
-    ) -> Sequence[float]:
-        assert self._raw_results is not None, "Expecting raw results from the database."
-        assert self._raw_results.get("distances") is not None
-        assert self.configs.query, "Expecting query messages in self.configs"
-        idx = self.configs.query.index(query_message)
-        dist = self._raw_results.get("distances")
-        if dist is None:  # pragma: nocover
-            raise ValueError("QueryResult should contain distances!")
-        else:
-            return list(-i for i in dist[idx])
+    async def compute_similarity(self, results: list[QueryResult]):
+        """
+        Do nothing, because the QueryResult objects already contain distances.
+        """
+        pass

src/vectorcode/subcommands/query/types.py

@@ -1,7 +1,7 @@
 import heapq
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Literal, Optional, Sequence, Union
+from typing import Literal, Union
 
 import numpy
 
@@ -22,8 +22,8 @@ class QueryResult:
 
     path: str
     chunk: Chunk
-    query: Sequence[str]
-    scores: Sequence[float]
+    query: tuple[str, ...]
+    scores: tuple[float, ...]
 
     @classmethod
     def merge(cls, *results: "QueryResult") -> "QueryResult":
@@ -47,14 +47,17 @@ def merge(cls, *results: "QueryResult") -> "QueryResult":
     def group(
         *results: "QueryResult",
         key: Union[Literal["path"], Literal["chunk"]] = "path",
-        top_k: Optional[int] = None,
+        top_k: int | Literal["auto"] | None = 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 == "auto":
+            top_k = int(numpy.mean(tuple(len(i) for i in grouped_result.values())))
+
         if top_k and top_k > 0:
             for group in grouped_result.keys():
                 grouped_result[group] = heapq.nlargest(top_k, grouped_result[group])