RerankerBase.rerank doesn't need the query messages as a parameter.

Author: Davidyz

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, query_chunks)
+    return reranker.rerank(results)
 
 
 async def build_query_results(

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

@@ -38,7 +38,5 @@ def create(cls, configs: Config, **kwargs: Any):
             raise
 
     @abstractmethod
-    def rerank(
-        self, results: QueryResult, query_chunks: list[str]
-    ) -> list[str]:  # pragma: nocover
+    def rerank(self, results: QueryResult) -> list[str]:  # pragma: nocover
         raise NotImplementedError

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

@@ -25,6 +25,9 @@ def __init__(
         **kwargs: Any,
     ):
         super().__init__(configs)
+        assert self.configs.query is not None, (
+            "'configs' should contain the query messages."
+        )
         from sentence_transformers import CrossEncoder
 
         if configs.reranker_params.get("model_name_or_path") is None:
@@ -36,17 +39,18 @@ def __init__(
             )
         self.model = CrossEncoder(**configs.reranker_params)
 
-    def rerank(self, results, query_chunks) -> list[str]:
-        self.query_chunks = query_chunks
+    def rerank(self, results) -> list[str]:
+        assert self.configs.query
+        query_chunks = self.configs.query
         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(self.query_chunks)):
+        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]
             ranks = self.model.rank(
-                self.query_chunks[query_chunk_idx], chunk_docs, apply_softmax=True
+                query_chunks[query_chunk_idx], chunk_docs, apply_softmax=True
             )
             for rank in ranks:
                 if QueryInclude.chunk in self.configs.include:

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

@@ -21,7 +21,7 @@ class NaiveReranker(RerankerBase):
     def __init__(self, configs: Config, **kwargs: Any):
         super().__init__(configs)
 
-    def rerank(self, results, query_chunks) -> list[str]:
+    def rerank(self, results) -> list[str]:
         assert results["metadatas"] is not None
         assert results["distances"] is not None
         documents: DefaultDict[str, list[float]] = defaultdict(list)

tests/subcommands/query/test_query.py

@@ -89,7 +89,7 @@ async def test_get_query_result_files(mock_collection, mock_config):
         # Check reranker was used correctly
         mock_get_reranker.assert_called_once_with(mock_config)
         mock_reranker_instance.rerank.assert_called_once_with(
-            mock_collection.query.return_value, mock_config.query
+            mock_collection.query.return_value
         )
 
         # Check the result

tests/subcommands/query/test_reranker.py

@@ -83,7 +83,7 @@ def __init__(self, configs, **kwargs):
 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, ["foo", "bar"])
+    result = reranker.rerank(query_result)
 
     # Check the result is a list of paths with correct length
     assert isinstance(result, list)
@@ -104,7 +104,7 @@ def test_naive_reranker_handles_none_path(config, query_result):
     ]
 
     reranker = NaiveReranker(config)
-    result = reranker.rerank(query_result_with_none, ["foo", "bar"])
+    result = reranker.rerank(query_result_with_none)
 
     # Check the None path was handled without errors
     assert isinstance(result, list)
@@ -121,6 +121,19 @@ def test_cross_encoder_reranker_initialization(mock_cross_encoder: MagicMock, co
     assert reranker.n_result == config.n_result
 
 
+@patch("sentence_transformers.CrossEncoder")
+def test_cross_encoder_reranker_initialization_fallback_model_name(
+    mock_cross_encoder: MagicMock, config
+):
+    expected_params = {"model_name_or_path": "cross-encoder/ms-marco-MiniLM-L-6-v2"}
+    config.reranker_params = {}
+    reranker = CrossEncoderReranker(config)
+
+    # Verify constructor was called with correct parameters
+    mock_cross_encoder.assert_called_once_with(**expected_params)
+    assert reranker.n_result == config.n_result
+
+
 @patch("sentence_transformers.CrossEncoder")
 def test_cross_encoder_reranker_rerank(
     mock_cross_encoder, config, query_result, query_chunks
@@ -137,7 +150,7 @@ def test_cross_encoder_reranker_rerank(
 
     reranker = CrossEncoderReranker(config)
 
-    result = reranker.rerank(query_result, query_chunks)
+    result = reranker.rerank(query_result)
 
     # Verify the model was called with correct parameters
     mock_model.rank.assert_called()
@@ -167,7 +180,7 @@ def test_naive_reranker_document_selection_logic(naive_reranker_conf):
     }
 
     reranker = NaiveReranker(naive_reranker_conf)
-    result = reranker.rerank(query_result, naive_reranker_conf.query)
+    result = reranker.rerank(query_result)
 
     # Check that files are included (exact order depends on implementation details)
     assert len(result) > 0
@@ -189,7 +202,7 @@ def test_naive_reranker_with_chunk_ids(naive_reranker_conf):
         ],
     }
     reranker = NaiveReranker(naive_reranker_conf)
-    result = reranker.rerank(query_result, naive_reranker_conf.query)
+    result = reranker.rerank(query_result)
 
     assert isinstance(result, list)
     assert len(result) <= naive_reranker_conf.n_result
@@ -224,26 +237,19 @@ def test_cross_encoder_reranker_with_chunk_ids(
             ],
             "documents": [["doc1", "doc2"], ["doc3", "doc4"]],
         },
-        config.query,
     )
 
     assert isinstance(result, list)
     assert all(isinstance(id, str) for id in result)
     assert all(id in ["id1", "id2", "id3", "id4"] for id in result)
 
 
-def test_get_reranker():
-    config = Config(reranker="NaiveReranker")
-    assert get_reranker(config).configs.reranker == "NaiveReranker"
+def test_get_reranker(config, naive_reranker_conf):
+    assert get_reranker(naive_reranker_conf).configs.reranker == "NaiveReranker"
 
-    config = Config(reranker="CrossEncoderReranker", reranker_params={"device": "cpu"})
     reranker = get_reranker(config)
     assert reranker.configs.reranker == "CrossEncoderReranker"
 
-    config = Config(
-        reranker="cross-encoder/ms-marco-MiniLM-L-6-v2",
-        reranker_params={"device": "cpu"},
-    )
     reranker = get_reranker(config)
     assert reranker.configs.reranker == "CrossEncoderReranker", (
         "configs.reranker should fallback to 'CrossEncoderReranker'"