diff --git a/AGENTS.md b/AGENTS.md index c32fe84..b9fa8a8 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -31,7 +31,7 @@ Let errors surface instead of swallowing them: degraded paths (a missing optiona Tests live in `test_.py` under `tests/unit/` or `tests/integration/`; integration suites drive the Typer CLI end to end, including porcelain output. Use fixtures/stubs instead of real APIs and keep provider and network interactions mocked so tests stay offline. Assert behavior or structured output rather than brittle console formatting. Pair each behavior change with happy-path and failure coverage, especially for invalid modes, empty queries, malformed config, stale or missing indexes, optional extras such as `flashrank`, and platform-specific shell behavior. The offline suite cannot catch real provider regressions: before opening a PR that touches providers, indexing/search, or the MCP server, run a real end-to-end check that exercises your change (for example, `python -m vexor index --path . --mode code` plus a search) using your configured API credentials, or the local provider if you have none, and note the result in the PR. ## Commit & Pull Request Guidelines -Commit subjects and PR titles follow [Conventional Commits](https://www.conventionalcommits.org): `type(scope): description`, e.g. `feat(mcp): add stdio server` or `fix(cache): handle locked index database`. Allowed types are `feat`, `fix`, `docs`, `style`, `refactor`, `perf`, `test`, `build`, `ci`, `chore`, and `revert`; the scope is optional but encouraged and should name the touched area (`mcp`, `cache`, `cli`, `gui`, `providers`, ...). Use the imperative mood, keep subjects under ~72 characters, and mark breaking changes with `!` (for example, `feat(config)!: drop legacy keys`). CI enforces the format on PR titles (`.github/workflows/conventional-commits.yml`); squash merges inherit the PR title as the commit subject, so a compliant title keeps `main` history clean. Name branches `type/short-slug` aligned with the commit type (e.g. `feat/hybrid-search`, `docs/roadmap-recover`); do all work on a branch and land it through a PR rather than committing directly to `main`. PRs should explain motivation, list the commands/tests exercised, link issues, and include screenshots for GUI changes or pasted terminal output for CLI output changes. Call out config-schema, cache-schema, Python API, bundled skill, or provider/reranker changes so reviewers can check compatibility. When behavior or contributor workflow changes, update whichever of `README.md`, docs, plugin metadata, and this guide the change affects. +Commit subjects and PR titles follow [Conventional Commits](https://www.conventionalcommits.org): `type(scope): description`, e.g. `feat(mcp): add stdio server` or `fix(cache): handle locked index database`. Allowed types are `feat`, `fix`, `docs`, `style`, `refactor`, `perf`, `test`, `build`, `ci`, `chore`, and `revert`; the scope is optional but encouraged and should name the touched area (`mcp`, `cache`, `cli`, `gui`, `providers`, ...). Use the imperative mood, keep subjects under ~72 characters, and mark breaking changes with `!` (for example, `feat(config)!: drop legacy keys`). CI enforces the format on PR titles (`.github/workflows/conventional-commits.yml`); squash merges inherit the PR title as the commit subject, so a compliant title keeps `main` history clean. Name branches `type/short-slug` aligned with the commit type (e.g. `feat/hybrid-search`, `docs/roadmap-recover`); do all work on a branch and land it through a PR rather than committing directly to `main`. PRs should explain motivation, list the commands/tests exercised, link issues, and include screenshots for GUI changes or pasted terminal output for CLI output changes. Call out config-schema, cache-schema, Python API, bundled skill, or provider/reranker changes so reviewers can check compatibility. When behavior or contributor workflow changes, update whichever of `README.md`, docs, plugin metadata, and this guide the change affects — the bundled skill (`plugins/vexor/skills/vexor-cli/SKILL.md`) is the easiest to forget. ## Security, Configuration & Maintenance Never commit API keys, private proxy endpoints, or local credentials; use `vexor config --set-api-key`, provider env vars, or ignored `.env` files. Persistent data lives under `~/.vexor`: config/cache data in the main directory, FlashRank assets under `~/.vexor/flashrank`, and local embedding models under `~/.vexor/models`. Sanitize filesystem paths, extension filters, and exclude patterns before using them. Treat extracted document text, embedding payloads, and remote rerank/provider responses as untrusted input before writing to disk or surfacing them elsewhere. Keep docs, bundled skills, and release metadata in sync when commands, packaging, or setup flows change. diff --git a/README.md b/README.md index 518b012..2cd357c 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,8 @@ When you remember what a file *does* but forget its name or location, Vexor find Designed for both humans and AI coding assistants, enabling semantic file discovery in autonomous agent workflows. +Optional hybrid retrieval (`vexor config --rerank hybrid`) fuses full-corpus BM25 with dense search, so exact identifiers can surface even when dense retrieval misses them. + Per-project `.vexorignore` files give you full gitignore-style control over what gets indexed, and you can opt in to per-project indexes by creating a `.vexor/` directory or running `vexor index --local`. ## Install diff --git a/docs/api/python.md b/docs/api/python.md index 28aa104..61d3b2c 100644 --- a/docs/api/python.md +++ b/docs/api/python.md @@ -140,7 +140,7 @@ The `config` payload (dict/JSON) supports: - `embedding_dimensions`: integer or null - `auto_index`: boolean - `local_cuda`: boolean (local provider only) -- `rerank`: `off`, `bm25`, `flashrank`, `remote` +- `rerank`: `off`, `bm25`, `flashrank`, `remote`, `hybrid` - `flashrank_model`: string or null - `remote_rerank`: object with `base_url`, `api_key`, `model` diff --git a/docs/cli.md b/docs/cli.md index 51b40bf..8047546 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -39,6 +39,10 @@ | `--no-cache` | In-memory only; do not read/write index cache | | `--local` | With `index`, create and use `/.vexor/index.db` | +Rerank values are `off`, `bm25`, `flashrank`, `remote`, and `hybrid`. +`hybrid` fuses full-corpus BM25 and dense rankings, allowing exact lexical +matches outside the dense candidate window to be retrieved. + Porcelain output fields: `rank`, `similarity`, `path`, `chunk_index`, `start_line`, `end_line`, `preview` (line fields are `-` when unavailable). diff --git a/docs/configuration.md b/docs/configuration.md index 192572b..f4f4de6 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -30,6 +30,7 @@ vexor config --set-update-check false # disable the daily update notice (d vexor config --rerank bm25 # optional BM25 rerank for top-k results vexor config --rerank flashrank # FlashRank rerank (requires optional extra) vexor config --rerank remote # remote rerank via HTTP endpoint +vexor config --rerank hybrid # full-corpus BM25 + dense rank fusion vexor config --set-flashrank-model ms-marco-MultiBERT-L-12 # multilingual model vexor config --set-flashrank-model # reset FlashRank model to default vexor config --clear-flashrank # remove cached FlashRank models @@ -65,10 +66,17 @@ CI), merged over `~/.vexor/config.json`. Credential fields inside Rerank reorders the semantic results with a secondary ranker. Candidate sizing uses `clamp(int(--top * 2), 20, 150)`. +`hybrid` is different from the candidate rerankers: it scores BM25 over every +indexed chunk and fuses that full lexical ranking with the full dense ranking +using reciprocal rank fusion. Hybrid result scores are normalized rank-fusion +scores in `[0, 1]`, not cosine similarities. A score of `1.0` means a chunk was +ranked first by both retrieval paths. + Recommended defaults: - Keep `off` unless you want extra precision. - Use `bm25` for lightweight lexical boosts; it is fast and lightweight. +- Use `hybrid` when exact identifiers and lexical-only matches matter. - BM25 uses a multilingual tokenizer (Bert pre-tokenizer), so it can handle CJK better. - Use `flashrank` for stronger reranking (requires `pip install "vexor[flashrank]"` and downloads a model to `~/.vexor/flashrank`). @@ -77,6 +85,10 @@ Recommended defaults: - For Chinese or multi-language content, set `--set-flashrank-model ms-marco-MultiBERT-L-12`. - If unset, FlashRank defaults to `ms-marco-TinyBERT-L-2-v2`. +This release raises the index cache schema version. Older index caches are +invalidated and rebuilt automatically; cached embeddings are shared and reused, +so the migration does not require re-embedding unchanged content. + ## Providers: Remote vs Local Vexor supports both remote API providers (`openai`, `gemini`, `voyageai`, diff --git a/docs/mcp.md b/docs/mcp.md index 86edd3a..484f26d 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -45,6 +45,10 @@ codex mcp add vexor -- vexor mcp } ``` +The `reranker` response field can be `"hybrid"`. In that case, each result's +`score` is a normalized reciprocal-rank-fusion score in `[0, 1]`, rather than a +cosine similarity. + If `vexor` is not on the client's PATH, use an absolute command path, or `python` with `"args": ["-m", "vexor", "mcp"]`. diff --git a/docs/roadmap.md b/docs/roadmap.md index 9442e76..62a0800 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -10,7 +10,8 @@ never leaving the machine. ## P0 — Agent-first distribution -- Hybrid retrieval as a first-class path: fuse BM25 and dense scores +- **Shipped behind `--rerank hybrid`:** hybrid retrieval as a first-class path + fuses BM25 and dense scores (e.g. reciprocal rank fusion) during search instead of offering BM25 only as an opt-in reranker. Dependencies (`rank-bm25`, `tokenizers`) are already present. Pure semantic search is weak on exact identifiers; @@ -22,7 +23,7 @@ never leaving the machine. - Full-corpus BM25 needs term statistics persisted alongside the index cache — rebuilding them per query is O(corpus). Design this together with the `vectors.npy`/memmap work in P1. - - Flip the ranking default only after the evaluation benchmark (next + - The default flip remains pending. Flip the ranking default only after the evaluation benchmark (next item) confirms hybrid beats dense-only, and call the change out in release notes since result ordering shifts for existing users. - Publish an evaluation: token cost + answer quality of agent+Vexor vs @@ -58,6 +59,15 @@ never leaving the machine. - Add AST-aware `code` mode chunking for Go and Rust (tree-sitter support). - Support `.vexorignore` for per-project ignore rules. - Project-level local cache (per-folder cache root override). +- Project-level config (`/.vexor/config.json`) for behavior-only + settings such as mode preferences, rerank, extensions, and exclude + patterns. Security constraint: repo files are attacker-controllable + input, so the loader must whitelist behavior fields and reject + credentials and endpoints (`api_key`, `base_url`, `remote_rerank`) with + an explicit error — stricter than the `VEXOR_CONFIG_JSON` env override, + which permits `base_url`. Reuse the guarded `config_from_json(base=...)` + merge; `vexor config --show` and `vexor doctor` should display each + field's origin (global vs project). - Additional embedding providers (Azure). - Evaluate an optional LLM reranker that reads a bounded set of retrieved candidates and judges their relevance to the query. Keep dense/BM25 diff --git a/plugins/vexor/skills/vexor-cli/SKILL.md b/plugins/vexor/skills/vexor-cli/SKILL.md index 47206ee..b0630e9 100644 --- a/plugins/vexor/skills/vexor-cli/SKILL.md +++ b/plugins/vexor/skills/vexor-cli/SKILL.md @@ -47,6 +47,7 @@ vexor "" [--path ] [--mode ] [--ext .py,.md] [--exclude-patte ## Troubleshooting +- Searching for an exact identifier (function/class/constant name) with weak results: suggest `vexor config --rerank hybrid` once — it fuses exact lexical matching with semantic search. - Need ignored or hidden files: add `--include-hidden` and/or `--no-respect-gitignore`. - Scriptable output: use `--format porcelain` (TSV) or `--format porcelain-z` (NUL-delimited). - Get detailed help: `vexor search --help`. diff --git a/scripts/eval_hybrid.py b/scripts/eval_hybrid.py new file mode 100644 index 0000000..2f65584 --- /dev/null +++ b/scripts/eval_hybrid.py @@ -0,0 +1,154 @@ +"""Evaluate dense, legacy BM25 reranking, and full-corpus hybrid retrieval.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Sequence + +from vexor import api +from vexor.config import DEFAULT_LOCAL_MODEL, load_config + +ARMS = ("off", "bm25", "hybrid") + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--path", type=Path, default=Path("."), help="Repository root") + parser.add_argument("--mode", default="auto") + parser.add_argument("--top", type=int, default=10) + parser.add_argument("--provider") + parser.add_argument("--model") + parser.add_argument( + "--queries", + type=Path, + default=Path(__file__).with_name("eval_queries.jsonl"), + ) + parser.add_argument("--json", action="store_true", dest="as_json") + parser.add_argument("--verbose", action="store_true") + return parser.parse_args() + + +def _load_queries(path: Path) -> list[dict[str, str]]: + queries: list[dict[str, str]] = [] + with path.open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + if not line.strip(): + continue + item = json.loads(line) + if not isinstance(item.get("query"), str) or not isinstance( + item.get("expected"), str + ): + raise ValueError(f"Invalid query record on line {line_number}") + queries.append(item) + return queries + + +def _relative_result_paths(response, root: Path) -> list[str]: + paths: list[str] = [] + for result in response.results: + try: + relative = result.path.resolve().relative_to(root) + except ValueError: + relative = result.path + paths.append(relative.as_posix()) + return paths + + +def _metrics(details: Sequence[dict[str, object]]) -> dict[str, float]: + count = len(details) + if not count: + return {"mrr_at_10": 0.0, "hit_at_1": 0.0, "hit_at_5": 0.0} + reciprocal_sum = 0.0 + hit_one = 0 + hit_five = 0 + for detail in details: + rank = detail["rank"] + if isinstance(rank, int) and rank <= 10: + reciprocal_sum += 1.0 / rank + hit_one += int(rank == 1) + hit_five += int(isinstance(rank, int) and rank <= 5) + return { + "mrr_at_10": reciprocal_sum / count, + "hit_at_1": hit_one / count, + "hit_at_5": hit_five / count, + } + + +def main() -> int: + args = _parse_args() + root = args.path.expanduser().resolve() + config = load_config() + provider = args.provider or config.provider + model = args.model or config.model + if provider == "local" and args.model is None: + model = DEFAULT_LOCAL_MODEL + common = { + "path": root, + "mode": args.mode, + "provider": provider, + "model": model, + "use_config": True, + } + api.index(**common) + + queries = _load_queries(args.queries) + by_arm: dict[str, list[dict[str, object]]] = {arm: [] for arm in ARMS} + for item in queries: + for arm in ARMS: + response = api.search( + item["query"], + top=args.top, + auto_index=False, + config={"rerank": arm}, + **common, + ) + returned = _relative_result_paths(response, root) + try: + rank: int | None = returned.index(item["expected"]) + 1 + except ValueError: + rank = None + by_arm[arm].append( + { + "query": item["query"], + "expected": item["expected"], + "rank": rank, + "results": returned, + } + ) + + output = { + "query_count": len(queries), + "top": args.top, + "arms": { + arm: {"metrics": _metrics(details), "queries": details} + for arm, details in by_arm.items() + }, + } + if args.as_json: + print(json.dumps(output, indent=2)) + return 0 + + print("| Rerank | MRR@10 | Hit@1 | Hit@5 |") + print("|---|---:|---:|---:|") + for arm in ARMS: + metrics = output["arms"][arm]["metrics"] + print( + f"| {arm} | {metrics['mrr_at_10']:.3f} | " + f"{metrics['hit_at_1']:.3f} | {metrics['hit_at_5']:.3f} |" + ) + if args.verbose: + for item in queries: + ranks = [] + for arm in ARMS: + detail = next( + row for row in by_arm[arm] if row["query"] == item["query"] + ) + ranks.append(f"{arm}={detail['rank'] or '-'}") + print(f"- {item['query']} -> {item['expected']} ({', '.join(ranks)})") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/eval_queries.jsonl b/scripts/eval_queries.jsonl new file mode 100644 index 0000000..5394489 --- /dev/null +++ b/scripts/eval_queries.jsonl @@ -0,0 +1,30 @@ +{"query": "_apply_bm25_rerank", "expected": "vexor/services/search_service.py"} +{"query": "load_index_vectors", "expected": "vexor/cache.py"} +{"query": "EMBED_CACHE_TTL_DAYS", "expected": "vexor/cache.py"} +{"query": "compare_snapshot", "expected": "vexor/cache.py"} +{"query": "BertPreTokenizer", "expected": "vexor/bm25.py"} +{"query": "_resolve_rerank_candidates", "expected": "vexor/services/search_service.py"} +{"query": "VEXOR_CONFIG_JSON", "expected": "vexor/config.py"} +{"query": "backfill_chunk_lines", "expected": "vexor/cache.py"} +{"query": "build_exclude_spec", "expected": "vexor/utils.py"} +{"query": "DEFAULT_FLASHRANK_MODEL", "expected": "vexor/config.py"} +{"query": "where is the sqlite schema created", "expected": "vexor/cache.py"} +{"query": "how does incremental indexing decide what to re-embed", "expected": "vexor/services/index_service.py"} +{"query": "render search results as a table", "expected": "vexor/cli.py"} +{"query": "load cached query embedding", "expected": "vexor/cache.py"} +{"query": "extract keywords from document content", "expected": "vexor/services/keyword_service.py"} +{"query": "validate embedding dimensions for a model", "expected": "vexor/config.py"} +{"query": "filter indexed results by extension", "expected": "vexor/services/search_service.py"} +{"query": "detect stale index cache", "expected": "vexor/services/cache_service.py"} +{"query": "configure an OpenAI compatible provider", "expected": "vexor/providers/openai.py"} +{"query": "start the MCP stdio server", "expected": "vexor/services/mcp_service.py"} +{"query": "mcp search tool schema", "expected": "vexor/services/mcp_service.py"} +{"query": "flashrank model cache directory", "expected": "vexor/config.py"} +{"query": "plugins/vexor skill installer", "expected": "vexor/services/skill_service.py"} +{"query": "docs configuration rerank strategies", "expected": "docs/configuration.md"} +{"query": "tests CLI doctor command", "expected": "tests/integration/test_cli.py"} +{"query": "local embedding provider ONNX", "expected": "vexor/providers/local.py"} +{"query": "JavaScript TypeScript parser chunks", "expected": "vexor/services/js_parser.py"} +{"query": "query cache sqlite table", "expected": "vexor/cache.py"} +{"query": "porcelain search output fields", "expected": "docs/cli.md"} +{"query": "roadmap hybrid retrieval benchmark", "expected": "docs/roadmap.md"} diff --git a/tests/integration/test_cli.py b/tests/integration/test_cli.py index d9b534b..df36a3d 100644 --- a/tests/integration/test_cli.py +++ b/tests/integration/test_cli.py @@ -82,6 +82,29 @@ def fake_perform_search(request): assert captured["mode"] == "auto" +def test_search_outputs_hybrid_reranker_label(tmp_path, monkeypatch): + sample_file = tmp_path / "exact.py" + sample_file.write_text("data") + + def fake_perform_search(_request): + return SearchResponse( + base_path=tmp_path, + backend="fake-backend", + results=[SearchResult(path=sample_file, score=1.0)], + is_stale=False, + index_empty=False, + reranker="hybrid", + ) + + monkeypatch.setattr("vexor.cli.perform_search", fake_perform_search) + result = CliRunner().invoke( + app, ["search", "exact", "--path", str(tmp_path)] + ) + + assert result.exit_code == 0 + assert "Reranker: hybrid" in result.stdout + + def test_search_no_respect_gitignore_flag_sets_false(tmp_path, monkeypatch): runner = CliRunner() sample_file = tmp_path / "alpha.txt" @@ -883,6 +906,19 @@ def test_config_set_and_show(tmp_path): assert "FlashRank model" not in strip_ansi(result_show.stdout) +def test_config_hybrid_rerank_round_trip(tmp_path): + runner = CliRunner() + result = runner.invoke(app, ["config", "--rerank", "hybrid"]) + + assert result.exit_code == 0 + config_path = tmp_path / "config" / "config.json" + assert json.loads(config_path.read_text())["rerank"] == "hybrid" + + shown = runner.invoke(app, ["config", "--show"]) + assert shown.exit_code == 0 + assert "hybrid" in shown.stdout + + def test_config_set_flashrank_model_empty_resets_to_default(tmp_path): runner = CliRunner() diff --git a/tests/unit/test_bm25.py b/tests/unit/test_bm25.py new file mode 100644 index 0000000..0205a42 --- /dev/null +++ b/tests/unit/test_bm25.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import math + +import numpy as np + +from vexor import bm25 + + +def test_tokenize_identifiers_cjk_and_punctuation() -> None: + identifier_tokens = bm25.tokenize("_apply_bm25_rerank") + assert {"apply", "bm25", "rerank"}.issubset(identifier_tokens) + assert "_apply_bm25_rerank" in identifier_tokens + assert bm25.tokenize("中文测试") + assert bm25.tokenize("!!!") == [] + + +def test_tokenize_plain_words_are_not_double_counted() -> None: + tokens = bm25.tokenize("plain config words") + + assert tokens.count("plain") == 1 + assert tokens.count("config") == 1 + assert tokens.count("words") == 1 + + +def test_term_frequencies() -> None: + assert bm25.term_frequencies(["alpha", "beta", "alpha"]) == { + "alpha": 2, + "beta": 1, + } + + +def test_score_postings_matches_hand_computed_values() -> None: + postings = { + "alpha": [(0, 2, 3), (1, 1, 2)], + "beta": [(1, 1, 2)], + } + scores = bm25.score_postings(["alpha", "beta"], postings, 3, 2.0) + + alpha_idf = math.log((3 - 2 + 0.5) / (2 + 0.5) + 1) + beta_idf = math.log((3 - 1 + 0.5) / (1 + 0.5) + 1) + expected_zero = alpha_idf * 2 * 2.5 / (2 + 1.5 * (0.25 + 0.75 * 3 / 2)) + expected_one = alpha_idf + beta_idf + assert math.isclose(scores[0], expected_zero, rel_tol=1e-12) + assert math.isclose(scores[1], expected_one, rel_tol=1e-12) + assert bm25.score_postings(["alpha"], postings, 3, 0.0) == {} + + +def test_rrf_fuse_normalizes_and_keeps_dense_only_rows() -> None: + fused = bm25.rrf_fuse([0, 1, 2], {0: 3.0, 2: 2.0}, 3) + + expected_dense_only = bm25.RRF_DENSE_WEIGHT * (bm25.RRF_K + 1) / ( + bm25.RRF_K + 2 + ) + expected_row_two = ( + bm25.RRF_DENSE_WEIGHT * (bm25.RRF_K + 1) / (bm25.RRF_K + 3) + + bm25.RRF_BM25_WEIGHT * (bm25.RRF_K + 1) / (bm25.RRF_K + 2) + ) + assert fused.dtype == np.float32 + assert fused[0] == 1.0 + assert math.isclose(fused[1], expected_dense_only, rel_tol=1e-6) + assert math.isclose(fused[2], expected_row_two, rel_tol=1e-6) + assert fused[2] > fused[1] > 0 + assert list(np.argsort(-fused, kind="stable")) == [0, 2, 1] + + +def test_rrf_fuse_dense_weight_prevents_bm25_backed_rows_displacing_dense_top() -> None: + dense_order = list(range(52)) + bm25_rows = range(47, 52) + bm25_scores = {row: float(52 - row) for row in bm25_rows} + + fused = bm25.rrf_fuse(dense_order, bm25_scores, 52) + + dense_top_expected = bm25.RRF_DENSE_WEIGHT + backed_expected = { + row: ( + bm25.RRF_DENSE_WEIGHT + * (bm25.RRF_K + 1) + / (bm25.RRF_K + row + 1) + + bm25.RRF_BM25_WEIGHT + * (bm25.RRF_K + 1) + / (bm25.RRF_K + bm25_rank) + ) + for bm25_rank, row in enumerate(bm25_rows, start=1) + } + equal_weight_scores = { + row: ( + 0.5 * (bm25.RRF_K + 1) / (bm25.RRF_K + row + 1) + + 0.5 * (bm25.RRF_K + 1) / (bm25.RRF_K + bm25_rank) + ) + for bm25_rank, row in enumerate(bm25_rows, start=1) + } + + assert math.isclose(fused[0], dense_top_expected, rel_tol=1e-6) + for row, expected in backed_expected.items(): + assert math.isclose(fused[row], expected, rel_tol=1e-6) + assert fused[row] < fused[0] + assert all(score > 0.5 for score in equal_weight_scores.values()) + assert int(np.argmax(fused)) == 0 diff --git a/tests/unit/test_cache.py b/tests/unit/test_cache.py index 514df78..0cf94f4 100644 --- a/tests/unit/test_cache.py +++ b/tests/unit/test_cache.py @@ -84,6 +84,58 @@ def test_store_and_load_index(tmp_path, monkeypatch): assert meta["extensions"] == () +def test_store_and_load_bm25_statistics(tmp_path, monkeypatch): + monkeypatch.setattr(cache, "CACHE_DIR", tmp_path / "cache") + root = tmp_path / "project" + root.mkdir() + file_path = root / "alpha.py" + file_path.write_text("alpha") + entry = cache.IndexedChunk( + path=file_path, + rel_path="alpha.py", + chunk_index=0, + preview="alpha", + embedding=[1.0, 0.0], + bm25_terms={"alpha": 2, "python": 1}, + bm25_doc_len=3, + ) + cache.store_index( + root=root, + model="model", + include_hidden=False, + mode=MODE, + recursive=True, + entries=[entry], + ) + _, _, metadata = cache.load_index_vectors( + root=root, + model="model", + include_hidden=False, + mode=MODE, + recursive=True, + ) + index_id = metadata["index_id"] + chunk_id = metadata["chunk_ids"][0] + + assert cache.load_bm25_stats(index_id) == (1, 3.0) + assert cache.load_bm25_postings(index_id, ["alpha", "missing"]) == { + "alpha": [(chunk_id, 2, 3)] + } + + +def test_bm25_readers_fall_back_when_tables_are_missing(tmp_path, monkeypatch): + monkeypatch.setattr(cache, "CACHE_DIR", tmp_path / "cache") + connection = cache._connect(cache.cache_db_path()) + try: + cache._ensure_schema(connection) + connection.execute("DROP TABLE bm25_posting") + connection.execute("DROP TABLE bm25_doc") + assert cache.load_bm25_stats(1, conn=connection) == (0, 0.0) + assert cache.load_bm25_postings(1, ["alpha"], conn=connection) == {} + finally: + connection.close() + + def test_cache_dir_context_overrides_cache_db_path(tmp_path): original_cache_dir = cache.CACHE_DIR with cache.cache_dir_context(tmp_path): @@ -458,6 +510,87 @@ def test_apply_index_updates_handles_add_modify_delete(tmp_path, monkeypatch): assert meta["dimension"] == 2 +def test_apply_index_updates_rewrites_cascades_and_preserves_bm25( + tmp_path, monkeypatch +): + monkeypatch.setattr(cache, "CACHE_DIR", tmp_path / "cache") + root = tmp_path / "project" + root.mkdir() + files = [root / name for name in ("a.txt", "b.txt", "c.txt")] + for file_path in files: + file_path.write_text(file_path.stem) + entries = [] + for idx, (file_path, term) in enumerate( + zip(files, ("alpha", "beta", "gamma")) + ): + entries.append( + cache.IndexedChunk( + path=file_path, + rel_path=file_path.name, + chunk_index=0, + preview=term, + embedding=np.array([1.0, float(idx)], dtype=np.float32), + bm25_terms={term: 1}, + bm25_doc_len=1, + ) + ) + cache.store_index( + root=root, + model="model", + include_hidden=False, + mode=MODE, + recursive=True, + entries=entries, + ) + _, _, before_meta = cache.load_index_vectors( + root, "model", False, MODE, True + ) + index_id = before_meta["index_id"] + gamma_before = cache.load_bm25_postings(index_id, ["gamma"]) + + files[0].write_text("delta") + stat_c = files[2].stat() + changed = cache.IndexedChunk( + path=files[0], + rel_path="a.txt", + chunk_index=0, + preview="delta", + embedding=np.array([0.0, 1.0], dtype=np.float32), + bm25_terms={"delta": 1}, + bm25_doc_len=1, + ) + cache.apply_index_updates( + root=root, + model="model", + include_hidden=False, + mode=MODE, + recursive=True, + ordered_entries=[("a.txt", 0), ("c.txt", 0)], + changed_entries=[changed], + touched_entries=[ + ( + "c.txt", + 0, + stat_c.st_size, + stat_c.st_mtime, + "gamma", + None, + None, + "", + ) + ], + removed_rel_paths=["b.txt"], + ) + + postings = cache.load_bm25_postings( + index_id, ["alpha", "beta", "gamma", "delta"] + ) + assert "alpha" not in postings + assert "beta" not in postings + assert postings["delta"][0][1:] == (1, 1) + assert postings["gamma"] == gamma_before["gamma"] + + def test_apply_index_updates_allows_deletions_without_embeddings(tmp_path, monkeypatch): monkeypatch.setattr(cache, "CACHE_DIR", tmp_path / "cache") root = tmp_path / "project" @@ -829,3 +962,40 @@ def test_project_cache_context_preserves_existing_gitignore(tmp_path, monkeypatc pass assert (marker / ".gitignore").read_text(encoding="utf-8") == "custom\n" + + +def test_compare_snapshot_current_for_nested_paths(tmp_path, monkeypatch): + """Stored posix rel_paths must match compare_snapshot's rel_path form.""" + monkeypatch.setattr(cache, "CACHE_DIR", tmp_path / "cache") + + root = tmp_path / "project" + sub = root / "sub" + sub.mkdir(parents=True) + nested = sub / "nested.txt" + nested.write_text("data", encoding="utf-8") + + entries = [ + cache.IndexedChunk( + path=nested, + rel_path="sub/nested.txt", + chunk_index=0, + preview="nested", + embedding=np.array([1.0, 0.0], dtype=np.float32), + ) + ] + cache.store_index( + root=root, + model="model", + include_hidden=False, + mode=MODE, + recursive=True, + entries=entries, + ) + meta = cache.load_index(root, "model", False, MODE, True) + + assert cache.compare_snapshot( + root, + False, + meta["files"], + recursive=True, + ) is True diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index c68023e..0da9fac 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -114,6 +114,14 @@ def test_save_and_load_rerank(tmp_path, monkeypatch): assert cfg.rerank == "bm25" +def test_hybrid_rerank_is_supported_and_round_trips(tmp_path, monkeypatch): + _prepare_config(tmp_path, monkeypatch) + config_module.save_config(config_module.Config(rerank="hybrid")) + + assert "hybrid" in config_module.SUPPORTED_RERANKERS + assert config_module.load_config().rerank == "hybrid" + + def test_save_and_load_flashrank_model(tmp_path, monkeypatch): _prepare_config(tmp_path, monkeypatch) diff --git a/tests/unit/test_init_service_extra.py b/tests/unit/test_init_service_extra.py index 84ba5f0..5eefbb6 100644 --- a/tests/unit/test_init_service_extra.py +++ b/tests/unit/test_init_service_extra.py @@ -155,6 +155,11 @@ def test_collect_rerank_settings_variants(monkeypatch): "remote_rerank_api_key": "secret", } + _prompt_sequence(monkeypatch, ["hybrid"]) + assert init_service._collect_rerank_settings(dry_run=False) == { + "rerank": "hybrid" + } + def test_alias_helpers_and_prompt_alias_setup(monkeypatch, tmp_path): monkeypatch.setenv("SHELL", "/bin/zsh") diff --git a/tests/unit/test_mcp_service.py b/tests/unit/test_mcp_service.py index 0a8a10d..97c1f29 100644 --- a/tests/unit/test_mcp_service.py +++ b/tests/unit/test_mcp_service.py @@ -239,6 +239,21 @@ def test_success_results_include_structured_content(tmp_path): assert result["structuredContent"]["results"][0]["rank"] == 1 +def test_search_response_carries_hybrid_reranker(tmp_path): + server, client = make_server(tmp_path) + original_search = client.search + + def hybrid_search(query, **kwargs): + response = original_search(query, **kwargs) + response.reranker = "hybrid" + return response + + client.search = hybrid_search + response = server.handle_message(tool_call(SEARCH_TOOL, {"query": "q"})) + + assert decode_tool_payload(response)["reranker"] == "hybrid" + + def test_error_results_omit_structured_content(tmp_path): server, _ = make_server(tmp_path, search_error=RuntimeError("boom")) response = server.handle_message(tool_call(SEARCH_TOOL, {"query": "q"})) diff --git a/tests/unit/test_search_service.py b/tests/unit/test_search_service.py index 3486aed..fe867d5 100644 --- a/tests/unit/test_search_service.py +++ b/tests/unit/test_search_service.py @@ -30,6 +30,161 @@ def test_bm25_tokenizer_handles_cjk() -> None: assert tokens +def _hybrid_request(tmp_path: Path, query: str, rerank: str) -> SearchRequest: + return SearchRequest( + query=query, + directory=tmp_path, + include_hidden=False, + respect_gitignore=True, + mode="name", + recursive=True, + top_k=1, + model_name="model", + batch_size=0, + provider="local", + base_url=None, + api_key=None, + local_cuda=False, + exclude_patterns=(), + extensions=(), + auto_index=False, + rerank=rerank, + ) + + +def test_hybrid_retrieves_lexical_match_outside_dense_candidate_clamp( + tmp_path: Path, +) -> None: + from vexor import bm25 + from vexor.services import search_service as service + + paths = [tmp_path / f"dense-{idx}.txt" for idx in range(24)] + paths.append(tmp_path / "lexical-only.txt") + vectors = np.array( + [[1.0 - idx * 0.01, 0.0] for idx in range(24)] + [[-1.0, 0.0]], + dtype=np.float32, + ) + chunks = [] + for idx, path in enumerate(paths): + document = "needle" if idx == 24 else "semantic candidate" + tokens = bm25.tokenize(document) + chunks.append( + { + "path": path.name, + "chunk_index": 0, + "preview": document, + "bm25_terms": bm25.term_frequencies(tokens), + "bm25_doc_len": len(tokens), + } + ) + metadata = {"chunks": chunks} + meta_getter = service._chunk_meta_from_entries(chunks) + + legacy, legacy_label = service._rank_results( + _hybrid_request(tmp_path, "needle", "bm25"), + paths=paths, + file_vectors=vectors, + query_vector=np.array([1.0, 0.0], dtype=np.float32), + chunk_meta_getter=meta_getter, + ) + hybrid, hybrid_label = service._rank_results( + _hybrid_request(tmp_path, "needle", "hybrid"), + paths=paths, + file_vectors=vectors, + query_vector=np.array([1.0, 0.0], dtype=np.float32), + chunk_meta_getter=meta_getter, + lexical_scorer=service._hybrid_scorer_from_entries(chunks), + ) + + assert legacy_label == "bm25" + assert legacy[0].path != paths[-1] + assert hybrid_label == "hybrid" + assert hybrid[0].path == paths[-1] + + +def test_hybrid_ranks_exact_identifier_above_sub_token_matches(tmp_path: Path) -> None: + from vexor import bm25 + from vexor.services import search_service as service + + documents = ["alpha beta gamma" for _ in range(12)] + documents.insert(1, "alpha_beta_gamma") + paths = [tmp_path / f"chunk-{idx}.txt" for idx in range(len(documents))] + vectors = np.array( + [[0.99 - idx * 0.01, 0.0] for idx in range(len(documents))], + dtype=np.float32, + ) + vectors[1] = [1.0, 0.0] + chunks = [] + for document in documents: + tokens = bm25.tokenize(document) + chunks.append( + { + "chunk_index": 0, + "preview": document, + "bm25_terms": bm25.term_frequencies(tokens), + "bm25_doc_len": len(tokens), + } + ) + + results, reranker = service._rank_results( + _hybrid_request(tmp_path, "alpha_beta_gamma", "hybrid"), + paths=paths, + file_vectors=vectors, + query_vector=np.array([1.0, 0.0], dtype=np.float32), + chunk_meta_getter=service._chunk_meta_from_entries(chunks), + lexical_scorer=service._hybrid_scorer_from_entries(chunks), + ) + + assert reranker == "hybrid" + assert results[0].path == paths[1] + + +def test_hybrid_empty_tokens_fall_back_to_dense(tmp_path: Path) -> None: + from vexor.services import search_service as service + + chunks = [ + { + "chunk_index": 0, + "bm25_terms": {"alpha": 1}, + "bm25_doc_len": 1, + } + ] + results, reranker = service._rank_results( + _hybrid_request(tmp_path, "!!!", "hybrid"), + paths=[tmp_path / "a.txt"], + file_vectors=np.array([[1.0, 0.0]], dtype=np.float32), + query_vector=np.array([1.0, 0.0], dtype=np.float32), + chunk_meta_getter=service._chunk_meta_from_entries(chunks), + lexical_scorer=service._hybrid_scorer_from_entries(chunks), + ) + + assert results[0].score == 1.0 + assert reranker is None + + +def test_hybrid_caps_unique_query_terms(tmp_path: Path) -> None: + from vexor.services import search_service as service + + captured: list[str] = [] + + def scorer(terms): + captured.extend(terms) + return {0: 1.0} + + scorer.has_data = True + query = " ".join(f"term{idx}" for idx in range(40)) + service._rank_results( + _hybrid_request(tmp_path, query, "hybrid"), + paths=[tmp_path / "a.txt"], + file_vectors=np.array([[1.0, 0.0]], dtype=np.float32), + query_vector=np.array([1.0, 0.0], dtype=np.float32), + chunk_meta_getter=service._chunk_meta_from_entries([{}]), + lexical_scorer=scorer, + ) + + assert len(captured) == 32 + + def test_perform_search_auto_indexes_when_missing(monkeypatch, tmp_path: Path) -> None: calls: dict[str, object] = {"load": 0, "indexed": 0, "index_kwargs": None} diff --git a/tests/unit/test_system_service.py b/tests/unit/test_system_service.py index 1f3590f..1cade43 100644 --- a/tests/unit/test_system_service.py +++ b/tests/unit/test_system_service.py @@ -143,6 +143,29 @@ def test_check_rerank_flashrank_ready(monkeypatch): assert result.passed is True +def test_check_rerank_hybrid_ready_and_degraded(monkeypatch): + monkeypatch.setattr(system_service.importlib.util, "find_spec", lambda _name: object()) + ready = system_service.check_rerank_configured( + "hybrid", + flashrank_model=None, + remote_rerank=None, + skip_api_test=False, + ) + assert ready is not None + assert ready.passed is True + + monkeypatch.setattr(system_service.importlib.util, "find_spec", lambda _name: None) + degraded = system_service.check_rerank_configured( + "hybrid", + flashrank_model=None, + remote_rerank=None, + skip_api_test=False, + ) + assert degraded is not None + assert degraded.passed is True + assert "degraded" in degraded.message.lower() + + def test_check_rerank_remote_incomplete(): result = system_service.check_rerank_configured( "remote", diff --git a/vexor/bm25.py b/vexor/bm25.py new file mode 100644 index 0000000..4142a15 --- /dev/null +++ b/vexor/bm25.py @@ -0,0 +1,115 @@ +"""Shared BM25 tokenization, scoring, and rank-fusion helpers.""" + +from __future__ import annotations + +from collections import Counter +from functools import lru_cache +import math +import re +from typing import Mapping, Sequence + +import numpy as np + +BM25_K1 = 1.5 +BM25_B = 0.75 +RRF_K = 60 +# Intentionally mirrors the legacy reranker's _FUSION_SEMANTIC_WEIGHT split. +RRF_DENSE_WEIGHT = 0.7 +RRF_BM25_WEIGHT = 0.3 +MAX_QUERY_TERMS = 32 + +_TOKEN_RE = re.compile(r"[A-Za-z0-9_]+") + + +@lru_cache(maxsize=1) +def _get_bm25_tokenizer(): + try: + from tokenizers.pre_tokenizers import BertPreTokenizer + except Exception: + return None + return BertPreTokenizer() + + +def tokenize(text: str) -> list[str]: + tokenizer = _get_bm25_tokenizer() + if tokenizer is None: + return _TOKEN_RE.findall(text.lower()) + tokens = [token for token, _ in tokenizer.pre_tokenize_str(text)] + normalized: list[str] = [] + for token in tokens: + cleaned = token.strip() + if not cleaned: + continue + if any(ch.isalnum() for ch in cleaned): + normalized.append(cleaned.lower()) + sub_tokens = set(normalized) + normalized.extend( + whole_token + for whole_token in _TOKEN_RE.findall(text.lower()) + if whole_token not in sub_tokens + ) + return normalized + + +def build_document(rel_path: str, label: str) -> str: + """Build the canonical lexical document for an indexed chunk.""" + + return f"{rel_path} {label}" + + +def term_frequencies(tokens: Sequence[str]) -> dict[str, int]: + return dict(Counter(tokens)) + + +def score_postings( + query_terms: Sequence[str], + postings: Mapping[str, Sequence[tuple[int, int, int]]], + doc_count: int, + avg_doc_len: float, +) -> dict[int, float]: + """Score matching posting lists with non-negative-idf Okapi BM25.""" + + if doc_count <= 0 or avg_doc_len <= 0: + return {} + scores: dict[int, float] = {} + for term in query_terms: + term_postings = postings.get(term, ()) + if not term_postings: + continue + df = len(term_postings) + idf = math.log((doc_count - df + 0.5) / (df + 0.5) + 1.0) + for chunk_id, tf, doc_len in term_postings: + denominator = tf + BM25_K1 * ( + 1.0 - BM25_B + BM25_B * doc_len / avg_doc_len + ) + if denominator <= 0: + continue + contribution = idf * tf * (BM25_K1 + 1.0) / denominator + scores[chunk_id] = scores.get(chunk_id, 0.0) + contribution + return scores + + +def rrf_fuse( + dense_order: Sequence[int], + bm25_scores_by_row: Mapping[int, float], + total_rows: int, + *, + k: int = RRF_K, +) -> np.ndarray: + """Fuse dense and BM25 rankings with normalized reciprocal rank fusion.""" + + fused = np.zeros(total_rows, dtype=np.float32) + for rank, row in enumerate(dense_order, start=1): + if 0 <= row < total_rows: + fused[row] += RRF_DENSE_WEIGHT * (k + 1.0) / (k + rank) + bm25_order = sorted( + ( + (row, score) + for row, score in bm25_scores_by_row.items() + if score > 0 and 0 <= row < total_rows + ), + key=lambda item: (-item[1], item[0]), + ) + for rank, (row, _score) in enumerate(bm25_order, start=1): + fused[row] += RRF_BM25_WEIGHT * (k + 1.0) / (k + rank) + return fused diff --git a/vexor/cache.py b/vexor/cache.py index cb9ac31..4f2d3a3 100644 --- a/vexor/cache.py +++ b/vexor/cache.py @@ -25,7 +25,7 @@ "vexor_cache_dir_override", default=None, ) -CACHE_VERSION = 6 +CACHE_VERSION = 7 DB_FILENAME = "index.db" EMBED_CACHE_TTL_DAYS = 30 EMBED_CACHE_MAX_ENTRIES = 50_000 @@ -47,6 +47,8 @@ class IndexedChunk: mtime: float | None = None start_line: int | None = None end_line: int | None = None + bm25_terms: Mapping[str, int] | None = None + bm25_doc_len: int | None = None def _cache_key( @@ -363,6 +365,8 @@ def _reset_index_schema(conn: sqlite3.Connection) -> None: DROP TABLE IF EXISTS file_embedding; DROP TABLE IF EXISTS chunk_embedding; DROP TABLE IF EXISTS chunk_meta; + DROP TABLE IF EXISTS bm25_posting; + DROP TABLE IF EXISTS bm25_doc; DROP TABLE IF EXISTS indexed_chunk; DROP TABLE IF EXISTS indexed_file; DROP TABLE IF EXISTS index_metadata; @@ -425,6 +429,19 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: end_line INTEGER ); + CREATE TABLE IF NOT EXISTS bm25_doc ( + chunk_id INTEGER PRIMARY KEY REFERENCES indexed_chunk(id) ON DELETE CASCADE, + token_count INTEGER NOT NULL + ); + + CREATE TABLE IF NOT EXISTS bm25_posting ( + index_id INTEGER NOT NULL REFERENCES index_metadata(id) ON DELETE CASCADE, + chunk_id INTEGER NOT NULL REFERENCES indexed_chunk(id) ON DELETE CASCADE, + term TEXT NOT NULL, + tf INTEGER NOT NULL, + PRIMARY KEY (index_id, term, chunk_id) + ) WITHOUT ROWID; + CREATE TABLE IF NOT EXISTS query_cache ( id INTEGER PRIMARY KEY AUTOINCREMENT, index_id INTEGER NOT NULL REFERENCES index_metadata(id) ON DELETE CASCADE, @@ -455,6 +472,9 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: CREATE INDEX IF NOT EXISTS idx_embedding_cache_lookup ON embedding_cache(model, text_hash); + + CREATE INDEX IF NOT EXISTS idx_bm25_posting_chunk + ON bm25_posting(chunk_id); """ ) @@ -623,6 +643,30 @@ def store_index( for idx, row in enumerate(inserted_ids) ), ) + bm25_doc_rows: list[tuple[int, int]] = [] + bm25_posting_rows: list[tuple[int, int, str, int]] = [] + for idx, row in enumerate(inserted_ids): + entry = entries[idx] + if entry.bm25_terms is None: + continue + chunk_id = int(row["id"]) + doc_len = int(entry.bm25_doc_len or 0) + bm25_doc_rows.append((chunk_id, doc_len)) + bm25_posting_rows.extend( + (int(index_id), chunk_id, term, int(tf)) + for term, tf in entry.bm25_terms.items() + ) + conn.executemany( + "INSERT INTO bm25_doc (chunk_id, token_count) VALUES (?, ?)", + bm25_doc_rows, + ) + conn.executemany( + """ + INSERT INTO bm25_posting (index_id, chunk_id, term, tf) + VALUES (?, ?, ?, ?) + """, + bm25_posting_rows, + ) conn.executemany( """ INSERT OR REPLACE INTO chunk_meta ( @@ -815,6 +859,31 @@ def apply_index_updates( for idx, row in enumerate(inserted_ids) ), ) + bm25_doc_rows: list[tuple[int, int]] = [] + bm25_posting_rows: list[tuple[int, int, str, int]] = [] + for idx, row in enumerate(inserted_ids): + entry = chunk_list[idx] + if entry.bm25_terms is None: + continue + chunk_id = int(row["id"]) + bm25_doc_rows.append( + (chunk_id, int(entry.bm25_doc_len or 0)) + ) + bm25_posting_rows.extend( + (int(index_id), chunk_id, term, int(tf)) + for term, tf in entry.bm25_terms.items() + ) + conn.executemany( + "INSERT INTO bm25_doc (chunk_id, token_count) VALUES (?, ?)", + bm25_doc_rows, + ) + conn.executemany( + """ + INSERT INTO bm25_posting (index_id, chunk_id, term, tf) + VALUES (?, ?, ?, ?) + """, + bm25_posting_rows, + ) conn.executemany( """ INSERT OR REPLACE INTO chunk_meta ( @@ -1383,6 +1452,89 @@ def load_chunk_metadata( connection.close() +def load_bm25_stats( + index_id: int, + conn: sqlite3.Connection | None = None, +) -> tuple[int, float]: + """Return the BM25 document count and average document length.""" + + db_path = cache_db_path() + owns_connection = conn is None + try: + connection = conn or _connect(db_path, readonly=True) + except sqlite3.OperationalError: + return 0, 0.0 + try: + try: + _ensure_schema_readonly( + connection, + tables=("indexed_chunk", "bm25_doc"), + ) + except sqlite3.OperationalError: + return 0, 0.0 + row = connection.execute( + """ + SELECT COUNT(*) AS doc_count, AVG(d.token_count) AS avg_doc_len + FROM bm25_doc AS d + JOIN indexed_chunk AS c ON c.id = d.chunk_id + WHERE c.index_id = ? + """, + (index_id,), + ).fetchone() + if row is None or not row["doc_count"]: + return 0, 0.0 + return int(row["doc_count"]), float(row["avg_doc_len"] or 0.0) + finally: + if owns_connection: + connection.close() + + +def load_bm25_postings( + index_id: int, + terms: Sequence[str], + conn: sqlite3.Connection | None = None, +) -> dict[str, list[tuple[int, int, int]]]: + """Load BM25 posting lists for the requested terms.""" + + unique_terms = list(dict.fromkeys(term for term in terms if term)) + if not unique_terms: + return {} + db_path = cache_db_path() + owns_connection = conn is None + try: + connection = conn or _connect(db_path, readonly=True) + except sqlite3.OperationalError: + return {} + try: + try: + _ensure_schema_readonly( + connection, + tables=("bm25_doc", "bm25_posting"), + ) + except sqlite3.OperationalError: + return {} + results: dict[str, list[tuple[int, int, int]]] = {} + for term_chunk in _chunk_values(unique_terms, 900): + placeholders = ", ".join("?" for _ in term_chunk) + rows = connection.execute( + f""" + SELECT p.term, p.chunk_id, p.tf, d.token_count + FROM bm25_posting AS p + JOIN bm25_doc AS d ON d.chunk_id = p.chunk_id + WHERE p.index_id = ? AND p.term IN ({placeholders}) + """, + (index_id, *term_chunk), + ).fetchall() + for row in rows: + results.setdefault(row["term"], []).append( + (int(row["chunk_id"]), int(row["tf"]), int(row["token_count"])) + ) + return results + finally: + if owns_connection: + connection.close() + + def load_query_vector( index_id: int, query_hash: str, @@ -1793,4 +1945,6 @@ def _relative_path(path: Path, root: Path) -> str: rel = path.relative_to(root) except ValueError: rel = path - return str(rel) + # Must match the posix form index_service._relative_to_root stores in the + # DB, or compare_snapshot marks nested paths stale on Windows forever. + return rel.as_posix() diff --git a/vexor/config.py b/vexor/config.py index 5627f78..c0e8e03 100644 --- a/vexor/config.py +++ b/vexor/config.py @@ -49,7 +49,13 @@ DEFAULT_RERANK = "off" DEFAULT_FLASHRANK_MODEL = "ms-marco-TinyBERT-L-2-v2" DEFAULT_FLASHRANK_MAX_LENGTH = 256 -SUPPORTED_RERANKERS: tuple[str, ...] = ("off", "bm25", "flashrank", "remote") +SUPPORTED_RERANKERS: tuple[str, ...] = ( + "off", + "bm25", + "flashrank", + "remote", + "hybrid", +) SUPPORTED_EXTRACT_BACKENDS: tuple[str, ...] = ("auto", "thread", "process") ENV_CONFIG_JSON = "VEXOR_CONFIG_JSON" ENV_NO_UPDATE_CHECK = "VEXOR_NO_UPDATE_CHECK" diff --git a/vexor/services/index_service.py b/vexor/services/index_service.py index 3020c22..6fc0cee 100644 --- a/vexor/services/index_service.py +++ b/vexor/services/index_service.py @@ -2,9 +2,10 @@ from __future__ import annotations +from collections import Counter +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor import itertools import os -from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor from dataclasses import dataclass, field from datetime import datetime, timezone from enum import Enum @@ -17,6 +18,7 @@ from .content_extract_service import TEXT_EXTENSIONS from .js_parser import JSTS_EXTENSIONS from ..cache import CACHE_VERSION, IndexedChunk, backfill_chunk_lines +from .. import bm25 from ..config import ( DEFAULT_EMBED_CONCURRENCY, DEFAULT_EXTRACT_BACKEND, @@ -502,6 +504,8 @@ def build_index_in_memory( "chunk_index": entry.chunk_index, "start_line": entry.start_line, "end_line": entry.end_line, + "bm25_terms": entry.bm25_terms, + "bm25_doc_len": entry.bm25_doc_len, } ) if rel_path not in file_snapshot: @@ -783,7 +787,7 @@ def _relative_to_root(path: Path, root: Path) -> str: rel = path.relative_to(root) except ValueError: rel = path - return str(rel) + return rel.as_posix() def _cached_chunk_map(chunk_entries: Sequence[dict]) -> dict[str, list[int]]: @@ -1019,10 +1023,13 @@ def _build_index_entries( from ..cache import embedding_cache_key # local import for idx, payload in enumerate(payloads): stat = _stat_for_path(payload.file, stat_cache) + rel_path = _relative_to_root(payload.file, root) + document = bm25.build_document(rel_path, payload.label) + tokens = bm25.tokenize(document) entries.append( IndexedChunk( path=payload.file, - rel_path=_relative_to_root(payload.file, root), + rel_path=rel_path, chunk_index=payload.chunk_index, preview=payload.preview or "", embedding=embeddings[idx], @@ -1031,6 +1038,8 @@ def _build_index_entries( mtime=stat.st_mtime, start_line=payload.start_line, end_line=payload.end_line, + bm25_terms=dict(Counter(tokens)), + bm25_doc_len=len(tokens), ) ) return entries diff --git a/vexor/services/init_service.py b/vexor/services/init_service.py index d8cfd9d..6e998d3 100644 --- a/vexor/services/init_service.py +++ b/vexor/services/init_service.py @@ -294,6 +294,11 @@ def _collect_rerank_settings(*, dry_run: bool) -> dict[str, object]: Messages.INIT_OPTION_RERANK_FLASHRANK_DESC, ) _print_option("4", Messages.INIT_OPTION_RERANK_REMOTE, Messages.INIT_OPTION_RERANK_REMOTE_DESC) + _print_option( + "5", + Messages.INIT_OPTION_RERANK_HYBRID, + Messages.INIT_OPTION_RERANK_HYBRID_DESC, + ) console.print() rerank_choice = _prompt_choice( Messages.INIT_PROMPT_RERANK, @@ -307,9 +312,11 @@ def _collect_rerank_settings(*, dry_run: bool) -> dict[str, object]: "flashrank": "flashrank", "4": "remote", "remote": "remote", + "5": "hybrid", + "hybrid": "hybrid", }, default="1", - allowed="1/2/3/4", + allowed="1/2/3/4/5", ) console.print() diff --git a/vexor/services/search_service.py b/vexor/services/search_service.py index 2269569..03c8803 100644 --- a/vexor/services/search_service.py +++ b/vexor/services/search_service.py @@ -6,12 +6,12 @@ from functools import lru_cache from pathlib import Path import json -import re import numpy as np -from typing import Sequence, TYPE_CHECKING +from typing import Callable, Sequence, TYPE_CHECKING from urllib import error as urlerror from urllib import request as urlrequest +from .. import bm25 from ..config import ( DEFAULT_EMBED_CONCURRENCY, DEFAULT_EXTRACT_BACKEND, @@ -69,22 +69,18 @@ class SearchResponse: reranker: str | None = None -_TOKEN_RE = re.compile(r"[A-Za-z0-9_]+") -_BM25_K1 = 1.5 -_BM25_B = 0.75 +_TOKEN_RE = bm25._TOKEN_RE +_BM25_K1 = bm25.BM25_K1 +_BM25_B = bm25.BM25_B _FUSION_SEMANTIC_WEIGHT = 0.7 -@lru_cache(maxsize=1) -def _get_bm25_tokenizer(): - try: - from tokenizers.pre_tokenizers import BertPreTokenizer - except Exception: - return None - return BertPreTokenizer() +_get_bm25_tokenizer = bm25._get_bm25_tokenizer def _bm25_tokenize(text: str) -> list[str]: + # Keep the local getter indirection for callers that monkeypatch this + # long-standing private compatibility surface. tokenizer = _get_bm25_tokenizer() if tokenizer is None: return _TOKEN_RE.findall(text.lower()) @@ -99,6 +95,65 @@ def _bm25_tokenize(text: str) -> list[str]: return normalized +def _hybrid_scorer_from_cache( + index_id: int | None, + chunk_ids: Sequence[int], +) -> Callable[[Sequence[str]], dict[int, float]] | None: + """Build a full-corpus lexical scorer backed by persisted postings.""" + + if index_id is None or not chunk_ids: + return None + from .. import cache + + row_by_chunk_id = {int(chunk_id): row for row, chunk_id in enumerate(chunk_ids)} + + def score(query_terms: Sequence[str]) -> dict[int, float]: + doc_count, avg_doc_len = cache.load_bm25_stats(int(index_id)) + score.has_data = doc_count > 0 + if not score.has_data: + return {} + postings = cache.load_bm25_postings(int(index_id), query_terms) + scores_by_chunk = bm25.score_postings( + query_terms, postings, doc_count, avg_doc_len + ) + return { + row_by_chunk_id[chunk_id]: value + for chunk_id, value in scores_by_chunk.items() + if chunk_id in row_by_chunk_id + } + + score.has_data = False + return score + + +def _hybrid_scorer_from_entries( + chunk_entries: Sequence[dict], +) -> Callable[[Sequence[str]], dict[int, float]] | None: + """Build a transient full-corpus lexical scorer from in-memory chunks.""" + + postings: dict[str, list[tuple[int, int, int]]] = {} + doc_lengths: list[int] = [] + for row, entry in enumerate(chunk_entries): + terms = entry.get("bm25_terms") + doc_len = entry.get("bm25_doc_len") + if terms is None or doc_len is None: + continue + length = int(doc_len) + doc_lengths.append(length) + for term, tf in terms.items(): + postings.setdefault(str(term), []).append((row, int(tf), length)) + if not doc_lengths: + return None + doc_count = len(doc_lengths) + avg_doc_len = sum(doc_lengths) / doc_count + + def score(query_terms: Sequence[str]) -> dict[int, float]: + return bm25.score_postings(query_terms, postings, doc_count, avg_doc_len) + + score.has_data = True + return score + + def _build_rerank_document(result: SearchResult) -> str: preview = result.preview or "" return f"{result.path.name} {result.path.as_posix()} {preview}".strip() @@ -502,6 +557,7 @@ def _rank_results( file_vectors: np.ndarray, query_vector: np.ndarray, chunk_meta_getter, + lexical_scorer: Callable[[Sequence[str]], dict[int, float]] | None = None, ) -> tuple[list, str | None]: """Score the index against the query, then rank and optionally rerank.""" from ..search import SearchResult # local import @@ -528,6 +584,37 @@ def _rank_results( ) similarities = np.asarray(file_vectors @ query_vector, dtype=np.float32) + if rerank == "hybrid": + query_terms = list(dict.fromkeys(bm25.tokenize(request.query)))[ + : bm25.MAX_QUERY_TERMS + ] + if query_terms and lexical_scorer is not None: + bm25_scores_by_row = lexical_scorer(query_terms) + if bm25_scores_by_row or getattr(lexical_scorer, "has_data", False): + dense_order = np.argsort(-similarities, kind="stable") + fused = bm25.rrf_fuse( + dense_order, bm25_scores_by_row, len(paths) + ) + top_indices = _top_indices(fused, request.top_k) + chunk_meta_for = chunk_meta_getter(top_indices) + scored: list[SearchResult] = [] + for idx in top_indices: + chunk_meta = chunk_meta_for(idx) or {} + start_line = chunk_meta.get("start_line") + end_line = chunk_meta.get("end_line") + scored.append( + SearchResult( + path=paths[idx], + score=float(fused[idx]), + preview=chunk_meta.get("preview"), + chunk_index=int(chunk_meta.get("chunk_index", 0)), + start_line=( + int(start_line) if start_line is not None else None + ), + end_line=int(end_line) if end_line is not None else None, + ) + ) + return scored, "hybrid" top_indices = _top_indices(similarities, candidate_count) chunk_meta_for = chunk_meta_getter(top_indices) scored: list[SearchResult] = [] @@ -740,6 +827,13 @@ def load_state() -> _IndexState: file_vectors=state.file_vectors, query_vector=query_vector, chunk_meta_getter=_chunk_meta_from_cache(state.chunk_ids, state.chunk_entries), + lexical_scorer=( + _hybrid_scorer_from_cache( + state.metadata.get("index_id"), state.chunk_ids + ) + if (request.rerank or "").strip().lower() == "hybrid" + else None + ), ) return SearchResponse( base_path=request.directory, @@ -772,6 +866,11 @@ def search_from_vectors( file_vectors=file_vectors, query_vector=query_vector, chunk_meta_getter=_chunk_meta_from_entries(metadata.get("chunks", [])), + lexical_scorer=( + _hybrid_scorer_from_entries(metadata.get("chunks", [])) + if (request.rerank or "").strip().lower() == "hybrid" + else None + ), ) return SearchResponse( base_path=request.directory, diff --git a/vexor/services/system_service.py b/vexor/services/system_service.py index 9b22a74..92d2ff3 100644 --- a/vexor/services/system_service.py +++ b/vexor/services/system_service.py @@ -263,6 +263,18 @@ def check_rerank_configured( passed=True, message=Messages.DOCTOR_RERANK_BM25_READY, ) + if normalized == "hybrid": + if importlib.util.find_spec("tokenizers") is None: + return DoctorCheckResult( + name="Rerank", + passed=True, + message=Messages.DOCTOR_RERANK_HYBRID_DEGRADED, + ) + return DoctorCheckResult( + name="Rerank", + passed=True, + message=Messages.DOCTOR_RERANK_HYBRID_READY, + ) if normalized == "flashrank": if importlib.util.find_spec("flashrank") is None: return DoctorCheckResult( diff --git a/vexor/text.py b/vexor/text.py index d9fc597..61cde89 100644 --- a/vexor/text.py +++ b/vexor/text.py @@ -139,7 +139,7 @@ class Messages: HELP_SET_UPDATE_CHECK = ( "Enable/disable the daily background update check (default: enabled)." ) - HELP_SET_RERANK = "Set the rerank strategy (off, bm25, flashrank, remote)." + HELP_SET_RERANK = "Set the rerank strategy (off, bm25, flashrank, remote, hybrid)." HELP_SET_FLASHRANK_MODEL = ( "Set the FlashRank model name (omit or empty resets to default)." ) @@ -278,11 +278,13 @@ class Messages: INIT_OPTION_RERANK_OFF_DESC = "default" INIT_OPTION_RERANK_BM25 = "BM25" INIT_OPTION_RERANK_BM25_DESC = "almost no extra latency" + INIT_OPTION_RERANK_HYBRID = "Hybrid" + INIT_OPTION_RERANK_HYBRID_DESC = "BM25 + dense fusion; best for exact identifiers" INIT_OPTION_RERANK_FLASHRANK = "FlashRank" INIT_OPTION_RERANK_FLASHRANK_DESC = "more accurate, slower" INIT_OPTION_RERANK_REMOTE = "Remote" INIT_OPTION_RERANK_REMOTE_DESC = "network-dependent" - INIT_PROMPT_RERANK = "Choose 1/2/3/4" + INIT_PROMPT_RERANK = "Choose 1/2/3/4/5" INIT_FLASHRANK_MISSING = "FlashRank is not installed." INIT_CONFIRM_INSTALL_FLASHRANK = "Install FlashRank now? (vexor[flashrank])" INIT_CONFIRM_FLASHRANK_DOWNLOAD = "Download the FlashRank model now? (recommended)" @@ -465,6 +467,10 @@ class Messages: DOCTOR_LOCAL_FAILED = "Local model test failed" DOCTOR_RERANK_BM25_READY = "BM25 rerank ready" DOCTOR_RERANK_BM25_MISSING = "BM25 rerank dependencies missing" + DOCTOR_RERANK_HYBRID_READY = "Hybrid retrieval ready" + DOCTOR_RERANK_HYBRID_DEGRADED = ( + "Hybrid retrieval ready with regex tokenization; CJK tokenization is degraded" + ) DOCTOR_RERANK_FLASHRANK_READY = "FlashRank rerank ready (model: {model})" DOCTOR_RERANK_FLASHRANK_MISSING = "FlashRank rerank not installed" DOCTOR_RERANK_REMOTE_READY = "Remote rerank configured (model: {model})"