feat(core): add optional cross-encoder reranker stage to search#1143
feat(core): add optional cross-encoder reranker stage to search#1143ironhalik wants to merge 4 commits into
Conversation
Rescore the top vector/hybrid retrieval candidates with a cross-encoder before the final top-k cut, targeting the rank-6-8 near-misses measured in the semantic embedding provider: a RerankProvider protocol, fastembed (local ONNX) and litellm (Cohere/Jina/etc.) providers, and a factory with a process-wide singleton cache. Wired into SearchRepositoryBase before the slice in _search_vector_only and _search_hybrid, so one insertion covers both SQLite and Postgres. Hybrid's internal vector leg passes _apply_rerank=False so reranking runs once on the fused result. Disabled by default (reranker_enabled=False) — no latency or model-download regression for existing users. Default local model is jinaai/jina-reranker-v1-tiny-en, which benchmarked higher recall@5 than ms-marco-MiniLM at equal-or-lower latency on LoCoMo. reranker_max_document_chars (default 0 = no cap) bounds cross-encoder input length to trade a little quality for latency on long notes; the model still truncates to its own token limit. Scoring & robustness (from adversarial review): - Providers emit [0,1] relevance: the fastembed cross-encoder sigmoid-squashes its logits so the public score stays bounded and comparable across backends; un-reranked tail rows are demoted below the reranked floor to keep the returned score column monotonic and avoid mixing scales. - Permanent faults surface, transient faults degrade: a reranker timeout/outage falls back to retrieval order (logged), but missing deps and provider-contract breaks (incomplete/misaligned response) raise RerankProviderContractError. The API router maps missing-deps -> 400 and contract breaks -> 502. - The candidate chunk pool over-fetches (RERANK_POOL_CHUNK_FANOUT) so a few large multi-chunk notes can't starve the deduped rerank window. - Skip the (possibly paid) rerank call for pages entirely past the pool. - Reranker cache key includes the resolved cache dir (parity with embeddings, avoids the basicmachines-co#741/basicmachines-co#872 shared-singleton-wrong-dir bug). - Config validation rejects reranker_enabled without semantic search, and the litellm provider paired with the default fastembed model (litellm can't route it), with a clear message to set an explicit provider/model. Closes the local cross-encoder half of basicmachines-co#618. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Michal Weinert <michal@weinert.io>
Account-wide search runs a separate search per project, each independently reranked, then merges and sorts by score. Per-project reranker scores aren't comparable across projects, so the merged order was miscalibrated. Rerank the merged candidate pool once (in _search_all_projects) to restore one consistent global ranking. No-op when reranking is disabled; transient reranker failures fall back to per-project order; permanent faults surface — matching the repository rerank path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Michal Weinert <michal@weinert.io>
…rerank Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Michal Weinert <michal@weinert.io>
63ef501 to
4a49544
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 63ef501943
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| scores[index] = float(item["relevance_score"]) | ||
| seen.add(index) |
There was a problem hiding this comment.
Reject duplicate rerank result indexes
When LiteLLM returns a malformed results list that repeats an index while still including every expected index (for example [0, 1, 0] for two documents), this loop overwrites the earlier score and seen still equals the full expected set, so the provider contract break is accepted and search is reranked with arbitrary duplicate data. The surrounding contract says each index must appear exactly once; fail fast when index in seen before assigning the score.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good catch — confirmed and fixed in e2cd4ac.
[0, 1, 0] for two documents did slip through: the range check passed each item, the duplicate index overwrote its own score (last-write-wins), and seen == {0, 1} still satisfied the coverage check, so the contract break was accepted. Added an explicit if index in seen: raise RerankProviderContractError(...) before assigning, plus a test_duplicate_index_raises regression test alongside the existing out-of-range / incomplete-coverage cases.
A LiteLLM rerank response that repeats an index while still covering every expected one (e.g. [0, 1, 0] for two documents) slipped past the coverage check: `seen` still equalled the full set, so the contract break was accepted and the duplicate silently overwrote the earlier score (last-write-wins). Fail fast on a repeated index to hold the "each index exactly once" contract. Addresses Codex review feedback on basicmachines-co#1143. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Michal Weinert <michal@weinert.io>
Hi,
I'm happily using basic-memory in my day to day workload, and wanted to chip in with an improvement - reranker.
The implementation was done by Opus 4.8, plus a review by Fable and GPT5.6 Terra.
The adversarial-review skill from the repo was also used.
I didn't do the full LoCoMo benchmark because I'm short on hardware right now. My homelab box offers ~6GB of memory bandwidth :)
I've been running this branch for couple of days now, and I'm pretty happy with the results.
Summary
Adds an optional cross-encoder reranker stage to vector/hybrid search. After
retrieval, the top candidates are rescored by a cross-encoder that reads the query
and document together, recovering relevant results that bi-encoder/FTS ranking
left just below the top-k cutoff.
Closes #950. Closes #618.
Reranking is off by default — existing users see zero change until they opt in.
When enabled it defaults to a local fastembed ONNX cross-encoder (no API, no
network), with an optional litellm provider for hosted rerankers (Cohere/Jina/
Voyage/Bedrock/self-hosted OpenAI-compatible endpoints).
Motivation
Per #950, roughly half of LoCoMo retrieval misses are ranking failures rather
than recall failures — the gold document is in the candidate set, just below the
cutoff — and there's ~20× latency headroom to spend on fixing it. A cross-encoder
is the standard fix: it re-scores a small candidate pool with full query↔document
cross-attention instead of a dot product.
What it does
RerankProviderprotocol mirroringEmbeddingProvider, with two impls:fastembed(default) — local ONNX cross-encoder, shares the embeddingmodel cache, squashes raw logits to
[0, 1]via sigmoid.litellm— routes to any litellm-supported rerank API viaprovider/modelstrings (e.g.
cohere/rerank-v3.5).search_repository_basefor vector and hybridretrieval only. Explicit
text/title/permalinksearches keep user-requestedranking semantics and never rerank.
search_all_projects): each project reranksindependently, so per-project scores aren't comparable across projects. A single
merged rerank over the best candidates from the pooled results restores one
calibrated global ordering.
Design decisions
the opt-in keeps the default install unchanged.
create_rerank_providerreturnsNonewhen disabled, so the hot path stays allocation-free.jinaai/jina-reranker-v1-tiny-en. Benchmarks higher recall@5than
ms-marco-MiniLMat equal-or-lower latency on LoCoMo. (litellm can't routethis local model name, so the litellm provider requires an explicit
provider/modeloverride — enforced by a config validator that fails fast.)
reranker_candidates(default 20) arerescored; deeper pages skip the provider entirely (never rescore rows we won't
return — matters for paid APIs).
[0, 1]relevance while theun-reranked tail still holds raw retrieval scores (up to ~1.3 fused). The tail is
demoted strictly below the reranked floor so a tail row can't numerically outrank a
reranked one — one comparable ordering without a second provider call.
retrieval order — an opt-in enhancement must never turn good results into a search
outage. Permanent faults (missing deps, a provider that breaks its response
contract — wrong score count, out-of-range/duplicate indices) surface loudly
(
RerankProviderContractError→ 502) rather than silently returning un-rerankedresults with no signal.
identity (provider, model, litellm routing, resolved cache dir) so the model loads
once. The cache dir is part of the key to avoid the FastEmbed cache defaults to /tmp/fastembed_cache — breaks MCP in sandboxed runtimes (Codex, ephemeral /tmp) #741/FastEmbed/ONNX embedding model is loaded multiple times per process, leaking ~25 GB of native memory #872 class of wrong-cache
bug the embedding factory already guards against.
Configuration
All under
BasicMemoryConfig, off by default:reranker_enabledfalsesemantic_search_enabledreranker_providerfastembedlitellmreranker_modeljinaai/jina-reranker-v1-tiny-enprovider/modelfor litellmreranker_candidates20reranker_max_document_chars0(full)reranker_api_base/reranker_api_keynullBenchmarks
No gold-labeled recall benchmark (LoCoMo) here — the test box can't run it end to
end. Instead, a latency + ranking-quality measurement on a real 100-note
ops corpus, hand-labeled (n=20, one correct note per query). Directional, not
statistically significant, but every change traces to a real query.
Setup:
fastembed+jinaai/jina-reranker-v1-tiny-en,candidates=13,max_document_chars=1000, hybrid search. Same query set run twice, onlyreranker_enabledflipped.Ranking quality (gold = the single correct note, matched by permalink):
Zero regressions — the cross-encoder never demoted a correct answer. It promoted
buried-but-retrieved notes to the top (e.g. audit-log-shipping 8→1, restic→PBS
migration 4→1, fs-freeze experiment 6→1). The remaining ties are queries where
the gold note wasn't in the retrieval pool at all — a recall-stage ceiling the
reranker can't fix by design.
Testing
New coverage across the pipeline, providers, factory, and cross-project merge:
unsupported-provider, disabled→None.
rejected; explicit litellm model accepted.
transient error, contract-error propagation, misaligned-score-count raises,
hybrid reranks exactly once.
demotion, transient degrade, permanent-fault surface, non-semantic search types
skipped.
just check+ targetedpytestgreen.