diff --git a/src/vectorcode/subcommands/query/__init__.py b/src/vectorcode/subcommands/query/__init__.py index fec70da4..22ab5abc 100644 --- a/src/vectorcode/subcommands/query/__init__.py +++ b/src/vectorcode/subcommands/query/__init__.py @@ -93,7 +93,7 @@ async def get_query_result_files( num_query = configs.n_result if QueryInclude.chunk in configs.include: if filter: - filter = {"$and": [filter.copy(), {"$gte": 0}]} + filter = {"$and": [filter.copy(), {"start": {"$gte": 0}}]} else: filter["start"] = {"$gte": 0} else: @@ -122,9 +122,8 @@ async def get_query_result_files( return [] reranker = get_reranker(configs) - return await reranker.rerank( - convert_query_results(chroma_query_results, configs.query) - ) + converted_results = convert_query_results(chroma_query_results, configs.query) + return await reranker.rerank(converted_results) async def build_query_results( diff --git a/tests/subcommands/query/test_query.py b/tests/subcommands/query/test_query.py index 67251e4c..06559836 100644 --- a/tests/subcommands/query/test_query.py +++ b/tests/subcommands/query/test_query.py @@ -238,7 +238,7 @@ async def test_get_query_result_chunks_with_query_exclude(mock_collection, mock_ mock_collection.query.assert_called_once() _, kwargs = mock_collection.query.call_args assert kwargs["where"] == { - "$and": [{"path": {"$nin": ["/excluded/path.py"]}}, {"$gte": 0}] + "$and": [{"path": {"$nin": ["/excluded/path.py"]}}, {"start": {"$gte": 0}}] }