make rerank and compute_similarity async.

Author: Zhe Yu

src/vectorcode/subcommands/query/__init__.py

@@ -68,7 +68,7 @@ async def get_query_result_files(
         return []
 
     reranker = get_reranker(configs)
-    return reranker.rerank(results)
+    return await reranker.rerank(results)
 
 
 async def build_query_results(

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

@@ -46,7 +46,7 @@ def create(cls, configs: Config, **kwargs: Any):
             raise
 
     @abstractmethod
-    def compute_similarity(
+    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,
@@ -61,7 +61,7 @@ def compute_similarity(
         """
         raise NotImplementedError
 
-    def rerank(self, results: QueryResult | dict) -> list[str]:
+    async def rerank(self, results: QueryResult | dict) -> list[str]:
         self._raw_results = cast(QueryResult, results)
         query_chunks = self.configs.query
         assert query_chunks
@@ -72,7 +72,9 @@ def rerank(self, results: QueryResult | dict) -> list[str]:
             chunk_ids = results["ids"][query_chunk_idx]
             chunk_metas = results["metadatas"][query_chunk_idx]
             chunk_docs = results["documents"][query_chunk_idx]
-            scores = self.compute_similarity(chunk_docs, query_chunks[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))

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

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from typing import Any
 
@@ -32,8 +33,8 @@ def __init__(
             )
         self.model = CrossEncoder(**configs.reranker_params)
 
-    def compute_similarity(self, results: list[str], query_message: str):
-        return list(
-            float(i)
-            for i in self.model.predict([(chunk, query_message) for chunk in results])
+    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)

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

@@ -17,7 +17,7 @@ class NaiveReranker(RerankerBase):
     def __init__(self, configs: Config, **kwargs: Any):
         super().__init__(configs)
 
-    def compute_similarity(
+    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."

tests/subcommands/query/test_query.py

@@ -64,11 +64,13 @@ async def test_get_query_result_files(mock_collection, mock_config):
     # Mock the reranker
     with patch("vectorcode.subcommands.query.get_reranker") as mock_get_reranker:
         mock_reranker_instance = MagicMock()
-        mock_reranker_instance.rerank.return_value = [
-            "file1.py",
-            "file2.py",
-            "file3.py",
-        ]
+        mock_reranker_instance.rerank = AsyncMock(
+            return_value=[
+                "file1.py",
+                "file2.py",
+                "file3.py",
+            ]
+        )
         mock_get_reranker.return_value = mock_reranker_instance
 
         # Call the function
@@ -103,7 +105,7 @@ async def test_get_query_result_files_include_chunk(mock_collection, mock_config
 
     with patch("vectorcode.subcommands.query.reranker.NaiveReranker") as MockReranker:
         mock_reranker_instance = MagicMock()
-        mock_reranker_instance.rerank.return_value = ["chunk1"]
+        mock_reranker_instance.rerank = AsyncMock(return_value=["chunk1"])
         MockReranker.return_value = mock_reranker_instance
 
         await get_query_result_files(mock_collection, mock_config)
@@ -188,7 +190,7 @@ async def test_get_query_result_files_with_query_exclude(mock_collection, mock_c
         mock_expand_path.return_value = "/excluded/path.py"
 
         mock_reranker_instance = MagicMock()
-        mock_reranker_instance.rerank.return_value = ["file1.py", "file2.py"]
+        mock_reranker_instance.rerank = AsyncMock(return_value=["file1.py", "file2.py"])
         MockReranker.return_value = mock_reranker_instance
 
         # Call the function
@@ -211,7 +213,7 @@ async def test_get_query_result_files_with_cross_encoder_reranker(
         "vectorcode.subcommands.query.reranker.CrossEncoderReranker"
     ) as MockCrossEncoder:
         mock_reranker_instance = MagicMock()
-        mock_reranker_instance.rerank.return_value = ["file1.py", "file2.py"]
+        mock_reranker_instance.rerank = AsyncMock(return_value=["file1.py", "file2.py"])
         MockCrossEncoder.return_value = mock_reranker_instance
 
         # Call the function
@@ -266,7 +268,7 @@ async def test_get_query_result_files_chunking(mock_collection, mock_config):
         MockChunker.return_value = mock_chunker_instance
 
         mock_reranker_instance = MagicMock()
-        mock_reranker_instance.rerank.return_value = ["file1.py", "file2.py"]
+        mock_reranker_instance.rerank = AsyncMock(return_value=["file1.py", "file2.py"])
         MockReranker.return_value = mock_reranker_instance
 
         # Call the function
@@ -300,7 +302,7 @@ async def test_get_query_result_files_multiple_queries(mock_collection, mock_con
         MockChunker.return_value = mock_chunker_instance
 
         mock_reranker_instance = MagicMock()
-        mock_reranker_instance.rerank.return_value = ["file1.py", "file2.py"]
+        mock_reranker_instance.rerank = AsyncMock(return_value=["file1.py", "file2.py"])
         MockReranker.return_value = mock_reranker_instance
 
         # Call the function

tests/subcommands/query/test_reranker.py

@@ -81,10 +81,11 @@ def __init__(self, configs, **kwargs):
         TestReranker.create(Config())
 
 
-def test_naive_reranker_rerank(naive_reranker_conf, query_result):
+@pytest.mark.asyncio
+async def test_naive_reranker_rerank(naive_reranker_conf, query_result):
     """Test basic reranking functionality of NaiveReranker"""
     reranker = NaiveReranker(naive_reranker_conf)
-    result = reranker.rerank(query_result)
+    result = await reranker.rerank(query_result)
 
     # Check the result is a list of paths with correct length
     assert isinstance(result, list)
@@ -117,10 +118,9 @@ def test_cross_encoder_reranker_initialization_fallback_model_name(
     assert reranker.n_result == config.n_result
 
 
+@pytest.mark.asyncio
 @patch("sentence_transformers.CrossEncoder")
-def test_cross_encoder_reranker_rerank(
-    mock_cross_encoder, config, query_result, query_chunks
-):
+async def test_cross_encoder_reranker_rerank(mock_cross_encoder, config, query_result):
     mock_model = MagicMock()
     mock_cross_encoder.return_value = mock_model
 
@@ -141,44 +141,49 @@ def test_cross_encoder_reranker_rerank(
     )
 
     reranker = CrossEncoderReranker(config)
-    result = reranker.rerank(query_result)
+    result = await reranker.rerank(query_result)
 
     # Result assertions
     assert isinstance(result, list)
     assert all(isinstance(path, str) for path in result)
     assert len(result) <= config.n_result
 
 
-def test_naive_reranker_document_selection_logic(naive_reranker_conf, query_result):
+@pytest.mark.asyncio
+async def test_naive_reranker_document_selection_logic(
+    naive_reranker_conf, query_result
+):
     """Test that NaiveReranker correctly selects documents based on distances"""
     # Create a query result with known distances
 
     reranker = NaiveReranker(naive_reranker_conf)
-    result = reranker.rerank(query_result)
+    result = await reranker.rerank(query_result)
 
     # Check that files are included (exact order depends on implementation details)
     assert len(result) > 0
     # Common files should be present
     assert "file2.py" in result or "file3.py" in result
 
 
-def test_naive_reranker_with_chunk_ids(naive_reranker_conf, query_result):
+@pytest.mark.asyncio
+async def test_naive_reranker_with_chunk_ids(naive_reranker_conf, query_result):
     """Test NaiveReranker returns chunk IDs when QueryInclude.chunk is set"""
     naive_reranker_conf.include.append(
         QueryInclude.chunk
     )  # Assuming QueryInclude.chunk would be "chunk"
 
     reranker = NaiveReranker(naive_reranker_conf)
-    result = reranker.rerank(query_result)
+    result = await reranker.rerank(query_result)
 
     assert isinstance(result, list)
     assert len(result) <= naive_reranker_conf.n_result
     assert all(isinstance(id, str) for id in result)
     assert all(id.startswith("id") for id in result)  # Verify IDs not paths
 
 
+@pytest.mark.asyncio
 @patch("sentence_transformers.CrossEncoder")
-def test_cross_encoder_reranker_with_chunk_ids(
+async def test_cross_encoder_reranker_with_chunk_ids(
     mock_cross_encoder, config, query_result
 ):
     """Test CrossEncoderReranker returns chunk IDs when QueryInclude.chunk is set"""
@@ -192,7 +197,7 @@ def test_cross_encoder_reranker_with_chunk_ids(
     config.include = {QueryInclude.chunk}
     reranker = CrossEncoderReranker(config)
 
-    result = reranker.rerank(query_result)
+    result = await reranker.rerank(query_result)
 
     mock_model.predict.assert_called()
     assert isinstance(result, list)
@@ -231,7 +236,7 @@ def test_add_reranker_success():
 
     @add_reranker
     class TestReranker(RerankerBase):
-        def compute_similarity(self, results, query_message):
+        async def compute_similarity(self, results, query_message):
             return []
 
     assert len(get_available_rerankers()) == original_count + 1
@@ -249,7 +254,7 @@ def test_add_reranker_duplicate():
     # First registration should succeed
     @add_reranker
     class TestReranker(RerankerBase):
-        def rerank(self, results, query_chunks):
+        async def compute_similarity(self, results, query_message):
             return []
 
     # Second registration should fail