|
52 | 52 | import yaml |
53 | 53 | from bm25 import load_or_build as bm25_load_or_build |
54 | 54 | from embeddings import load_or_build as embeddings_load_or_build |
| 55 | +from reranker import rerank as _cross_encoder_rerank |
55 | 56 |
|
56 | 57 | from .paths import home as _agent_mem_home |
57 | 58 |
|
@@ -317,12 +318,17 @@ def _embedding_search( |
317 | 318 | log.exception("priming: embedding search raised") |
318 | 319 | return [] |
319 | 320 |
|
320 | | - # Filter out very weak semantic matches — cosine below ~0.25 is |
321 | | - # noise from this model and would poison RRF with irrelevant |
322 | | - # paths. Empirical: gibberish queries against unrelated entries |
323 | | - # produce scores in the 0.02-0.05 range; genuine paraphrase |
324 | | - # matches sit at 0.4+. |
325 | | - EMBEDDING_NOISE_FLOOR = 0.25 |
| 321 | + # Filter out very weak semantic matches before they reach RRF. |
| 322 | + # Calibrated for nomic-embed-text-v1.5 — its cosine distribution |
| 323 | + # sits noticeably higher than older models (unrelated ~0.40-0.50, |
| 324 | + # genuine paraphrase matches ~0.65+). A floor at 0.50 separates the |
| 325 | + # populations: nonsense queries against unrelated entries fall under |
| 326 | + # it (so gibberish in == gibberish-free out, the "no real match |
| 327 | + # clears stale priming" semantic), while real paraphrase matches |
| 328 | + # comfortably clear it. Re-tune if the embedder swaps again — see |
| 329 | + # test_search_returns_empty_for_unrelated_query in tools/search for |
| 330 | + # the empirical handle. |
| 331 | + EMBEDDING_NOISE_FLOOR = 0.50 |
326 | 332 | return [(Path(h.path), float(h.score)) for h in hits if float(h.score) >= EMBEDDING_NOISE_FLOOR] |
327 | 333 |
|
328 | 334 |
|
@@ -353,31 +359,100 @@ def _rrf_merge( |
353 | 359 | return ordered[:k_top] |
354 | 360 |
|
355 | 361 |
|
| 362 | +# Minimum cross-encoder logit for a candidate to count as a real match. |
| 363 | +# ms-marco-MiniLM-L-12-v2 outputs unnormalized relevance logits — strong |
| 364 | +# matches sit at +5 to +10, weak-but-real around -2 to +2, definite |
| 365 | +# non-matches at -5 and below. The "no real match" filter is mostly |
| 366 | +# handled upstream by ``EMBEDDING_NOISE_FLOOR``; this floor is the |
| 367 | +# rerank's own guard against an outright non-match slipping through |
| 368 | +# the embedding lane on lexical-only signal. A permissive value lets |
| 369 | +# weak-but-real matches through (a tightly-scoped lesson on a less |
| 370 | +# common topic can rerank around 0). Tighten toward 0 if noise leaks |
| 371 | +# through; loosen toward -5 if valid matches get dropped. |
| 372 | +_RERANK_SCORE_FLOOR = -3.0 |
| 373 | + |
| 374 | + |
| 375 | +def _rerank_candidates( |
| 376 | + query: str, |
| 377 | + candidates: List[Tuple[Path, float]], |
| 378 | + *, |
| 379 | + k: int, |
| 380 | +) -> List[Tuple[Path, float]]: |
| 381 | + """Cross-encoder rerank pass on ``candidates``; degrade to input on failure. |
| 382 | +
|
| 383 | + Reads each candidate's body once, hands the ``(query, body)`` pairs to |
| 384 | + the cross-encoder, and returns the top-``k`` by relevance score with |
| 385 | + the cross-encoder score in place of the RRF score. Candidates scoring |
| 386 | + below ``_RERANK_SCORE_FLOOR`` are dropped — see the constant's docstring |
| 387 | + for why. The reranker returns ``None`` on any internal failure (model |
| 388 | + missing, predict raised, empty input) — in that case we hand back the |
| 389 | + upstream order truncated to ``k`` so retrieval still works without |
| 390 | + rerank. |
| 391 | + """ |
| 392 | + if not candidates: |
| 393 | + return [] |
| 394 | + |
| 395 | + bodies: List[Tuple[Path, str]] = [] |
| 396 | + for path, _ in candidates: |
| 397 | + try: |
| 398 | + bodies.append((path, path.read_text(encoding="utf-8"))) |
| 399 | + except (OSError, UnicodeDecodeError): |
| 400 | + # Skip unreadable candidates rather than fail the whole rerank. |
| 401 | + continue |
| 402 | + if not bodies: |
| 403 | + return candidates[:k] |
| 404 | + |
| 405 | + reranked = _cross_encoder_rerank(query, bodies) |
| 406 | + if reranked is None: |
| 407 | + return candidates[:k] |
| 408 | + above_floor = [(p, s) for p, s in reranked if s >= _RERANK_SCORE_FLOOR] |
| 409 | + return above_floor[:k] |
| 410 | + |
| 411 | + |
356 | 412 | def _hybrid_search( |
357 | 413 | knowledge_dir: Path, |
358 | 414 | query: str, |
359 | 415 | *, |
360 | 416 | k: int, |
361 | 417 | ) -> List[Tuple[Path, float]]: |
362 | | - """Run BM25 and embedding search in parallel, merge via RRF. |
363 | | -
|
364 | | - Each lane pulls a few more than ``k`` so RRF has overlap to work |
365 | | - with. If the embedding lane is unavailable (no model installed, |
366 | | - fresh corpus, transient failure), gracefully degrades to BM25-only. |
| 418 | + """Run BM25 + embedding search, fuse with RRF, then cross-encoder rerank. |
| 419 | +
|
| 420 | + Three-stage retrieval, low precision → high precision: |
| 421 | +
|
| 422 | + 1. **Recall**: BM25 and embedding lanes each pull ``k * 4`` candidates. |
| 423 | + BM25 catches exact-token matches; embeddings catch paraphrase. |
| 424 | + 2. **Fusion**: RRF merges the two ranked lists into ``k * 4`` |
| 425 | + topically-relevant candidates (rank-based, robust to score scale |
| 426 | + differences between BM25 and cosine). |
| 427 | + 3. **Rerank**: a cross-encoder co-attends over each ``(query, body)`` |
| 428 | + pair and scores actual applicability — orthogonal signal to the |
| 429 | + overlap-based lanes upstream. Returns the top ``k`` reranked. |
| 430 | +
|
| 431 | + Graceful degradation: if the embedding lane is unavailable, we |
| 432 | + fall back to BM25-only feed into the reranker. If the reranker is |
| 433 | + unavailable, we fall back to the RRF order verbatim. The hot path |
| 434 | + must keep working when any single component breaks. |
367 | 435 | """ |
368 | | - # Pull 2× k from each lane so RRF has a chance to surface entries |
369 | | - # that one lane ranks low but the other ranks high. |
| 436 | + # Pull 2× k from each lane. Wider over-fetch (e.g. 4×) would give |
| 437 | + # the reranker more variety, but cross-encoder predict on MPS scales |
| 438 | + # nonlinearly — measured ~60ms @ 5 candidates, ~170ms @ 10, ~600ms @ |
| 439 | + # 20 (graph-cache effects above the small-tensor threshold). At 4× |
| 440 | + # the rerank stage alone overruns the 200ms hook budget. Stick to 2× |
| 441 | + # so end-to-end fits. |
370 | 442 | per_lane_k = max(k * 2, k + 5) |
371 | 443 | bm25_hits = _bm25_search(knowledge_dir, query, k=per_lane_k) |
372 | 444 | emb_hits = _embedding_search(knowledge_dir, query, k=per_lane_k) |
373 | 445 | if not bm25_hits and not emb_hits: |
374 | 446 | return [] |
| 447 | + |
375 | 448 | if not emb_hits: |
376 | | - # Embedding lane unavailable — fall back to BM25 verbatim. |
377 | | - return bm25_hits[:k] |
378 | | - if not bm25_hits: |
379 | | - return emb_hits[:k] |
380 | | - return _rrf_merge([bm25_hits, emb_hits], k_top=k) |
| 449 | + fused = bm25_hits |
| 450 | + elif not bm25_hits: |
| 451 | + fused = emb_hits |
| 452 | + else: |
| 453 | + fused = _rrf_merge([bm25_hits, emb_hits], k_top=per_lane_k) |
| 454 | + |
| 455 | + return _rerank_candidates(query, fused, k=k) |
381 | 456 |
|
382 | 457 |
|
383 | 458 | # Scope bonuses are calibrated against the RRF top-rank score (~0.016 for |
|
0 commit comments