Skip to content

Commit e76caee

Browse files
committed
refactor(cli)!: drop legacy reranker configs.
1 parent 7ad9d2a commit e76caee

File tree

3 files changed

+26
-17
lines changed

3 files changed

+26
-17
lines changed

src/vectorcode/subcommands/query/__init__.py

+9-2
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@
1414
get_collection,
1515
verify_ef,
1616
)
17-
from vectorcode.subcommands.query.reranker import get_reranker
17+
from vectorcode.subcommands.query.reranker import (
18+
RerankerError,
19+
get_reranker,
20+
)
1821

1922
logger = logging.getLogger(name=__name__)
2023

@@ -170,7 +173,11 @@ async def query(configs: Config) -> int:
170173
)
171174
configs.include = [QueryInclude.path, QueryInclude.document]
172175

173-
structured_result = await build_query_results(collection, configs)
176+
try:
177+
structured_result = await build_query_results(collection, configs)
178+
except RerankerError: # pragma: nocover
179+
# error logs should be handled where they're raised
180+
return 1
174181

175182
if configs.pipe:
176183
print(json.dumps(structured_result))

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

+10-5
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,14 @@
1818
}
1919

2020

21+
class RerankerError(Exception):
22+
pass
23+
24+
25+
class RerankerInitialisationError(RerankerError):
26+
pass
27+
28+
2129
def add_reranker(cls):
2230
"""
2331
This is a class decorator that allows you to add a custom reranker that can be
@@ -69,8 +77,5 @@ def get_reranker(configs: Config) -> RerankerBase:
6977
if not configs.reranker:
7078
return NaiveReranker(configs)
7179
else:
72-
configs.reranker_params.update({"model_name_or_path": configs.reranker})
73-
configs.reranker = "CrossEncoderReranker"
74-
return CrossEncoderReranker(
75-
configs,
76-
)
80+
logger.error(f"{configs.reranker} is not a valid reranker type!")
81+
raise RerankerInitialisationError()

tests/subcommands/query/test_query.py

+7-10
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
get_query_result_files,
1313
query,
1414
)
15+
from vectorcode.subcommands.query.reranker import (
16+
RerankerError,
17+
)
1518

1619

1720
@pytest.fixture
@@ -203,9 +206,7 @@ async def test_get_query_result_files_with_query_exclude(mock_collection, mock_c
203206

204207

205208
@pytest.mark.asyncio
206-
async def test_get_query_result_files_with_cross_encoder_reranker(
207-
mock_collection, mock_config
208-
):
209+
async def test_get_query_reranker_initialisation_error(mock_collection, mock_config):
209210
# Configure to use CrossEncoder reranker
210211
mock_config.reranker = "cross-encoder/model-name"
211212

@@ -216,13 +217,9 @@ async def test_get_query_result_files_with_cross_encoder_reranker(
216217
mock_reranker_instance.rerank = AsyncMock(return_value=["file1.py", "file2.py"])
217218
MockCrossEncoder.return_value = mock_reranker_instance
218219

219-
# Call the function
220-
result = await get_query_result_files(mock_collection, mock_config)
221-
222-
# Check CrossEncoder was used correctly
223-
MockCrossEncoder.assert_called_once()
224-
# Check the result
225-
assert result == ["file1.py", "file2.py"]
220+
with pytest.raises(RerankerError):
221+
# Call the function
222+
await get_query_result_files(mock_collection, mock_config)
226223

227224

228225
@pytest.mark.asyncio

0 commit comments

Comments
 (0)