diff --git a/ir/cli.py b/ir/cli.py index 14faaa3..3bd176f 100644 --- a/ir/cli.py +++ b/ir/cli.py @@ -53,13 +53,16 @@ def build(name, *, embedder=None, full=True): return f"built {name!r}: {len(corpus)} records (embedder {corpus.embedder_id})" -def search(name, query, *, k=10): - """Search a built corpus and print the top-k hits.""" +def search(name, query, *, k=10, mode="dense"): + """Search a built corpus and print the top-k hits. + + mode: dense (cosine) | lexical (BM25) | hybrid (dense + BM25 fused via RRF). + """ corpus = open_corpus(name) if len(corpus) == 0: return f"corpus {name!r} is empty; build it first: ir build {name}" lines = [] - for h in corpus.search(query, k=k): + for h in corpus.search(query, k=k, mode=mode): # artifact_id is the unique label across corpora (skill name[@parent], # package name, or a report's relative path). lines.append(f"{h.score:+.3f} {h.artifact_id} [{h.surface_kind}]") diff --git a/ir/retrieve.py b/ir/retrieve.py index 3b2db49..bb3099a 100644 --- a/ir/retrieve.py +++ b/ir/retrieve.py @@ -1,25 +1,46 @@ -"""Retrieval — hard metadata filtering + dense brute-force ranking. +"""Retrieval — hard metadata filtering + dense / lexical / hybrid ranking. :func:`search` embeds the query, applies a hard metadata filter to narrow the candidate set (the ``vd`` Mongo-style filter language — ownership, name, tags), -ranks the survivors by cosine similarity (exact brute force; correct and -instant at ``ir``'s corpus sizes), and collapses surface-hits to the best -surface per artifact. +then ranks the survivors. Three ranking modes share that one filtered candidate +set: -Hybrid lexical fusion (BM25 + ``vd.reciprocal_rank_fusion``) and cross-encoder -reranking are deliberate next-step seams (see the retrieval issue): dense + -filter is the shippable baseline. +- ``"dense"`` (default) — exact brute-force cosine over the embedding matrix. +- ``"lexical"`` — Okapi BM25 over the candidates' text (``vd.bm25_lexical_search``). +- ``"hybrid"`` — dense and lexical fused by Reciprocal Rank Fusion + (``vd.reciprocal_rank_fusion``), the rank-based fuse that sidesteps the + cosine/BM25 score-scale mismatch. + +Hybrid matters for short, identifier-heavy capability text (skill / package / +tool names), the regime where dense-only retrieval fails silently on exact +identifiers and rare terms. An optional ``rerank`` hook (any +:class:`ef.Reranker`) re-scores the top fused candidates with a cross-encoder; +it defaults to ``None`` so retrieval stays offline and API-free out of the box. + +All ranking is exact brute force — correct and instant at ``ir``'s corpus +sizes — and surface-hits collapse to the best surface per artifact. Lexical and +fusion reuse ``vd``; if ``vd`` is unavailable, hybrid degrades to dense and +lexical returns no results (both with a warning), so a missing optional dep +never hard-fails a search. """ from __future__ import annotations +import warnings from collections.abc import Iterable, Mapping -from typing import Any +from types import SimpleNamespace +from typing import Any, Callable import numpy as np from .base import SearchHit, best_per_artifact +#: Ranking modes accepted by :func:`search`. +MODES = ("dense", "lexical", "hybrid") + +#: RRF rank constant — the ``k`` of ``1 / (k + rank)`` (standard default 60). +DFLT_RRF_K = 60 + def _matches(metadata: Mapping[str, Any], filter: Mapping[str, Any]) -> bool: """Mongo-style metadata match, via ``vd`` when available.""" @@ -40,6 +61,139 @@ def _embed_query(embedder, query: str) -> np.ndarray: return vec / norm if norm else vec +def _filter_candidates( + metas: list[dict], + surfaces: Iterable[str] | None, + filter: Mapping[str, Any] | None, +) -> list[int]: + """Row indices surviving the surface restriction and hard metadata filter.""" + surface_set = set(surfaces) if surfaces is not None else None + candidates = [] + for j, m in enumerate(metas): + if surface_set is not None and m["surface_kind"] not in surface_set: + continue + if filter is not None and not _matches(m.get("metadata", {}), filter): + continue + candidates.append(j) + return candidates + + +def _dense_ranked( + corpus, + mat: np.ndarray, + ids: list[str], + candidates: list[int], + query: str, + fetch: int, +) -> list[tuple[str, float]]: + """Top-*fetch* ``(record_id, cosine)`` pairs by exact brute-force cosine.""" + qv = _embed_query(corpus.embedder, query) + if qv.shape[0] != mat.shape[1]: + raise ValueError( + f"Query embedding dim {qv.shape[0]} != index dim {mat.shape[1]}. " + f"The corpus was built with a different embedder than the one " + f"querying it; rebuild the corpus or use its original embedder." + ) + sub = mat[candidates] + scores = sub @ qv + order = np.argsort(-scores)[:fetch] + return [(ids[candidates[o]], float(scores[o])) for o in order] + + +def _lexical_ranked( + ids: list[str], + metas: list[dict], + candidates: list[int], + query: str, + fetch: int, + bm25: Mapping[str, Any] | None, +) -> list[tuple[str, float]]: + """Top-*fetch* ``(record_id, bm25)`` pairs via ``vd.bm25_lexical_search``. + + The candidate texts are exposed to ``vd`` as a zero-copy mapping view + (``record_id -> obj`` with ``.text`` / ``.metadata``) so no vectors are + duplicated and BM25 runs only over the already-filtered candidates. Returns + ``[]`` (with a warning) if ``vd``'s lexical search is unavailable, letting + hybrid degrade to dense. + """ + try: + from vd import bm25_lexical_search + except Exception: + warnings.warn( + "vd.bm25_lexical_search unavailable; BM25 lexical ranking is " + "skipped. Hybrid falls back to dense; lexical mode returns no " + "results. Install `vd` for lexical/hybrid retrieval.", + stacklevel=3, + ) + return [] + + collection = { + ids[j]: SimpleNamespace( + text=metas[j]["text"], metadata=metas[j].get("metadata", {}) + ) + for j in candidates + } + results = bm25_lexical_search(collection, query, limit=fetch, **(bm25 or {})) + return [(r["id"], float(r["score"])) for r in results] + + +def _rrf_fuse( + dense: list[tuple[str, float]], + lexical: list[tuple[str, float]], + rrf_k: int, + fetch: int, +) -> list[tuple[str, float]]: + """Fuse dense + lexical rankings via ``vd.reciprocal_rank_fusion``. + + Falls back to the dense ranking if ``vd`` is unavailable or lexical is + empty (RRF of a single non-empty list is just that list's order). + """ + if not lexical: + return dense[:fetch] + try: + from vd import reciprocal_rank_fusion + except Exception: + warnings.warn( + "vd.reciprocal_rank_fusion unavailable; falling back to dense " + "ranking for hybrid search.", + stacklevel=3, + ) + return dense[:fetch] + + fused = reciprocal_rank_fusion( + [ + [{"id": rid, "score": s} for rid, s in dense], + [{"id": rid, "score": s} for rid, s in lexical], + ], + k=rrf_k, + ) + return [(item["id"], float(item["rrf_score"])) for item in fused[:fetch]] + + +def _apply_rerank( + query: str, + ranked: list[tuple[str, float]], + meta_by_id: Mapping[str, dict], + rerank: Callable, + limit: int, +) -> list[tuple[str, float]]: + """Re-score the ranked candidates with an :class:`ef.Reranker` via ``ef.rerank``. + + The full candidate list is handed to the reranker; it may reorder freely + and the top ``limit`` are kept. The reranker score replaces the candidate + score, and ``_rid`` rides along on each segment so the reordered output + maps back to records. + """ + from ef import rerank as ef_rerank + + segments = [{"text": meta_by_id[rid]["text"], "_rid": rid} for rid, _ in ranked] + reordered = ef_rerank(query, segments, rerank, limit=limit) + return [ + (s["_rid"], float(s.get("metadata", {}).get("rerank_score", 0.0))) + for s in reordered + ] + + def search( corpus, query: str, @@ -48,6 +202,11 @@ def search( filter: Mapping[str, Any] | None = None, surfaces: Iterable[str] | None = None, per_artifact: bool = True, + mode: str = "dense", + rrf_k: int = DFLT_RRF_K, + fetch_k: int | None = None, + rerank: Callable | None = None, + bm25: Mapping[str, Any] | None = None, ) -> list[SearchHit]: """Return the top-*k* :class:`~ir.base.SearchHit` for *query*. @@ -57,45 +216,57 @@ def search( filter : a ``vd`` Mongo-style filter over record metadata (hard filter). surfaces : restrict to these surface kinds (e.g. ``{"description"}``). per_artifact : collapse to the best surface per artifact (default True). + mode : ``"dense"`` (default, cosine), ``"lexical"`` (BM25), or ``"hybrid"`` + (dense + BM25 fused by Reciprocal Rank Fusion). Hybrid is the strongest + default for short, identifier-heavy text; ``"dense"`` is the historical + behavior and is kept as the default for backward compatibility. + rrf_k : the RRF rank constant for ``"hybrid"`` (standard default 60). + fetch_k : candidate depth before fusion / reranking / dedupe + (default ``max(k*5, 50)`` when collapsing per artifact, else ``k``). + rerank : an optional :class:`ef.Reranker` (``(query, segments) -> scores``) + that re-scores the ranked candidates after fusion. The reranker sees + the full candidate list (up to ``fetch_k`` items) and may reorder it. + Default ``None`` keeps retrieval offline (no model download / API call). + bm25 : optional Okapi params forwarded to ``vd.bm25_lexical_search`` + (e.g. ``{"k1": 1.5, "b": 0.75}``). """ + if mode not in MODES: + raise ValueError(f"unknown mode {mode!r}; expected one of {MODES}") + ids, mat, metas = corpus.store.matrix() if not ids: return [] - surface_set = set(surfaces) if surfaces is not None else None - candidates = [] - for j, m in enumerate(metas): - if surface_set is not None and m["surface_kind"] not in surface_set: - continue - if filter is not None and not _matches(m.get("metadata", {}), filter): - continue - candidates.append(j) + candidates = _filter_candidates(metas, surfaces, filter) if not candidates: return [] - qv = _embed_query(corpus.embedder, query) - if qv.shape[0] != mat.shape[1]: - raise ValueError( - f"Query embedding dim {qv.shape[0]} != index dim {mat.shape[1]}. " - f"The corpus was built with a different embedder than the one " - f"querying it; rebuild the corpus or use its original embedder." - ) - sub = mat[candidates] - scores = sub @ qv + fetch = fetch_k if fetch_k is not None else (max(k * 5, 50) if per_artifact else k) + + if mode == "lexical": + ranked = _lexical_ranked(ids, metas, candidates, query, fetch, bm25) + else: + dense = _dense_ranked(corpus, mat, ids, candidates, query, fetch) + if mode == "dense": + ranked = dense + else: # hybrid + lexical = _lexical_ranked(ids, metas, candidates, query, fetch, bm25) + ranked = _rrf_fuse(dense, lexical, rrf_k, fetch) + + meta_by_id = {ids[j]: metas[j] for j in candidates} - # Over-fetch before collapsing to artifacts so dedupe doesn't starve top-k. - take = min(len(candidates), max(k * 5, 50) if per_artifact else k) - order = np.argsort(-scores)[:take] + if rerank is not None and ranked: + ranked = _apply_rerank(query, ranked, meta_by_id, rerank, fetch) hits = [ SearchHit( - artifact_id=metas[candidates[o]]["artifact_id"], - surface_kind=metas[candidates[o]]["surface_kind"], - score=float(scores[o]), - text=metas[candidates[o]]["text"], - metadata=metas[candidates[o]].get("metadata", {}), + artifact_id=meta_by_id[rid]["artifact_id"], + surface_kind=meta_by_id[rid]["surface_kind"], + score=score, + text=meta_by_id[rid]["text"], + metadata=meta_by_id[rid].get("metadata", {}), ) - for o in order + for rid, score in ranked ] if per_artifact: hits = best_per_artifact(hits) diff --git a/tests/test_hybrid.py b/tests/test_hybrid.py new file mode 100644 index 0000000..eaea72a --- /dev/null +++ b/tests/test_hybrid.py @@ -0,0 +1,104 @@ +"""Hybrid retrieval tests — dense + BM25 fused via RRF, plus the rerank hook. + +Hermetic: the light (numpy-only) hashing embedder, no model download, no +network. The core claim under test is that lexical (BM25) signal recovers an +*exact rare identifier* that dense-only retrieval misses, and that hybrid +fusion lifts that doc's rank relative to dense. +""" + +import pytest + +import ir +from ir.store import CorpusStore + +# `svc` is the only doc carrying the rare identifier ``zxqv``; the d* docs flood +# the corpus with the common query terms (deploy/systemd/unit), so their IDF is +# low and the frequency-weighted hashing embedder is distracted by d3's repeats. +DOCS = { + "svc": "zxqv service config", + "d1": "deploy systemd unit alpha", + "d2": "deploy systemd unit beta", + "d3": "deploy systemd unit gamma deploy systemd unit", + "d4": "deploy systemd unit delta", +} +QUERY = "deploy systemd unit zxqv" + + +def _corpus(): + src = ir.CorpusSource.from_mapping(DOCS, name="hyb", strategy=ir.WholeText()) + return ir.build(src, store=CorpusStore.memory(), embedder="light") + + +def _order(corpus, **kw): + return [h.artifact_id for h in ir.search(corpus, QUERY, k=10, **kw)] + + +def test_lexical_finds_rare_identifier(): + # BM25 ranks the rare-identifier doc first; dense buries it. + corpus = _corpus() + assert _order(corpus, mode="lexical")[0] == "svc" + assert _order(corpus, mode="dense")[0] != "svc" + + +def test_hybrid_lifts_what_dense_misses(): + corpus = _corpus() + dense_order = _order(corpus, mode="dense") + hybrid_order = _order(corpus, mode="hybrid") + # Hybrid recovers the lexically-strong doc: its rank improves vs dense. + assert hybrid_order.index("svc") < dense_order.index("svc") + + +def test_dense_is_the_default(): + # Omitting mode must equal mode="dense" (backward compatibility). + corpus = _corpus() + assert _order(corpus) == _order(corpus, mode="dense") + + +def test_unknown_mode_raises(): + corpus = _corpus() + with pytest.raises(ValueError): + ir.search(corpus, QUERY, mode="bogus") + + +def test_rerank_hook_reorders(): + # A reranker that prefers a specific token must reorder the results. + docs = {"a": "red apple fruit", "b": "blue sky cloud", "c": "green grass field"} + src = ir.CorpusSource.from_mapping(docs, name="rr", strategy=ir.WholeText()) + corpus = ir.build(src, store=CorpusStore.memory(), embedder="light") + + def prefer_grass(query, segments): + return [10.0 if "grass" in s["text"] else 0.0 for s in segments] + + reranked = [ + h.artifact_id for h in ir.search(corpus, "apple", k=3, rerank=prefer_grass) + ] + assert reranked[0] == "c" + + +def test_bm25_params_forwarded(): + # Passing Okapi params must not break lexical search (smoke test). + corpus = _corpus() + hits = ir.search(corpus, QUERY, mode="lexical", bm25={"k1": 1.2, "b": 0.5}) + assert hits and hits[0].artifact_id == "svc" + + +def _prefer_svc(query, segments): + # Reranker that promotes the doc whose text mentions "service". + return [10.0 if "service" in s["text"] else 0.0 for s in segments] + + +def test_rerank_applies_across_modes(): + # The rerank hook must run after ranking in every mode, not just dense. + corpus = _corpus() + for mode in ("dense", "lexical", "hybrid"): + top = ir.search(corpus, QUERY, k=5, mode=mode, rerank=_prefer_svc)[0] + assert top.artifact_id == "svc", f"rerank did not take effect for {mode=}" + + +def test_rerank_with_per_artifact_false(): + # rerank must also work when surfaces are not collapsed per artifact. + corpus = _corpus() + hits = ir.search( + corpus, QUERY, k=5, mode="hybrid", per_artifact=False, rerank=_prefer_svc + ) + assert hits and hits[0].artifact_id == "svc"