feat(rerank): wire endpoint-backed reranking into BM25 retrieval surfaces

Reuse the shipped Cohere/Jina rerank client as an optional post-process on
the BM25 surfaces (tool search, skill search, memory composition) via one
seam: BM25Index gains an injected reranker + a two-stage search (BM25 recall
top-50 -> rerank -> top-k). No new storage.

Gated on a configured endpoint plus tools.rerank_bm25 (default on, matching
rerank_web_search). tools.rerank_bm25_threshold (default 0.0 = off) is a
relevance FLOOR for proactive memory surfacing: BM25 always returns something,
so without a floor every-turn memory injection spends tokens on the top-k of
whatever lexically matched; the reranker score is what makes a meaningful
"inject nothing" gate possible.

Two reranker modes (BM25Index rerank_filters):
- REORDER (reactive tool/skill search): the reranker must never drop results
  -> fall back to BM25 order on empty, backfill omitted pool items, so a
  misbehaving endpoint can't silently lose tools.
- FILTER (memory, rerank_filters = threshold > 0): a clean empty/short result
  is honoured (inject nothing) -- a deliberate divergence from
  web_search._rerank_results.
Parse/endpoint failure is a discrete branch from the floor: an empty result
for non-empty input means an unparseable response (a conforming reranker
scores every doc), so the closure raises RerankError and BM25Index falls back
to BM25 order in BOTH modes -- the floor only acts on valid scores.

Also: cap the rerank client timeout at 15s (the per-turn memory path can't
afford tools.timeout's 120s default); move the Reranker alias to rerank.py
(shared, no import cycle); document the endpoint egress in the rerank_bm25
help, the admin Reranker-role description, and docs/tools.md; add
scripts/bench_bm25_rerank.py (manual, needs a live endpoint) to measure
precision@k/MRR lift and recommend a threshold default.

Negative-tested: reorder fallback-on-empty and omitted-item backfill,
filter-mode honor-empty, singleton-still-floored, the parse-fail RerankError
raise, the >= floor boundary, and pool-position-to-doc-index mapping -- each
guard reverted to confirm its test fails, then restored.
This commit is contained in:
Patrick Buckley
2026-06-01 12:40:39 -07:00
parent c6e0293185
commit 215f7506ba
15 changed files with 894 additions and 24 deletions
+2
View File
@@ -321,6 +321,8 @@ Search the web using a text query.
The `rerank_web_search` toggle defaults on once an endpoint is configured. If the endpoint is unreachable or errors, web_search falls back silently to the backend's native result order — reranking never makes a search fail.
When `rerank_bm25` is enabled, the candidate text for memory, tool, and skill retrieval (memory name/description/content and tool/skill names + descriptions) is also sent to the rerank endpoint — a self-hosted endpoint (vLLM/TEI/llama.cpp) keeps it on your infrastructure, a hosted provider (Cohere/Jina/Voyage) sends it off-box.
Example (vLLM serving a Qwen3 reranker):
```toml
+260
View File
@@ -0,0 +1,260 @@
#!/usr/bin/env python3
"""Manual BM25-vs-BM25→rerank benchmark for the retrieval surfaces.
Run this by hand on a host that has a Cohere/Jina-compatible rerank endpoint
configured (vLLM, TEI, llama.cpp, or a hosted Cohere/Jina/Voyage key). It is NOT
a pytest test and is never collected by the test suite — it needs a live endpoint
to do anything useful.
It compares plain BM25 top-k against the two-stage BM25→rerank path on two small
in-file labeled corpora (tool-like docs and synthetic memory dicts), printing
precision@k, MRR, a ranking diff, and the relevant-vs-irrelevant score
distribution so you can pick a sensible ``tools.rerank_bm25_threshold`` default.
The endpoint bearer token (for hosted providers) is read from the
``$TURNSTONE_RERANK_API_KEY`` environment variable, never a flag, so it does not
land in shell history or the process listing.
Example::
TURNSTONE_RERANK_API_KEY=... \
.venv/bin/python scripts/bench_bm25_rerank.py \
--rerank-url http://localhost:8000/rerank \
--rerank-model BAAI/bge-reranker-v2-m3 --k 3 --threshold 0.0
"""
from __future__ import annotations
import argparse
import os
import statistics
from typing import TYPE_CHECKING
from turnstone.core.bm25 import BM25Index
from turnstone.core.rerank import resolve_rerank_client
if TYPE_CHECKING:
from collections.abc import Callable
from turnstone.core.rerank import RerankClient
# (query, relevant-doc-name) labels over a handful of tool-like docs.
TOOL_DOCS: list[dict[str, str]] = [
{"name": "read_file", "description": "Read the contents of a file from disk"},
{"name": "write_file", "description": "Write or overwrite a file on disk"},
{"name": "list_dir", "description": "List the entries in a directory"},
{"name": "web_search", "description": "Search the web and return result snippets"},
{"name": "send_email", "description": "Send an email message to a recipient"},
{"name": "run_bash", "description": "Execute a shell command and capture output"},
{"name": "create_issue", "description": "Open a new issue in the bug tracker"},
{"name": "query_db", "description": "Run a SQL query against the database"},
]
TOOL_LABELS: list[tuple[str, str]] = [
("open a file and show me what's inside", "read_file"),
("look something up on the internet", "web_search"),
("send a message by email", "send_email"),
("run a terminal command", "run_bash"),
("file a bug report", "create_issue"),
("fetch rows from the database", "query_db"),
]
# Synthetic memory dicts (name + description + content) with labeled queries.
MEMORY_DOCS: list[dict[str, str]] = [
{
"name": "postgres_conn",
"description": "database connection settings",
"content": "host=db.internal port=5432 user=app sslmode=require",
},
{
"name": "redis_cache",
"description": "cache server config",
"content": "host=redis port=6379 maxmemory 2gb eviction allkeys-lru",
},
{
"name": "deploy_runbook",
"description": "production deploy steps",
"content": "build image, push to registry, roll nodes one at a time",
},
{
"name": "oncall_rotation",
"description": "who is on call",
"content": "primary alice, secondary bob, escalate to carol after 30m",
},
{
"name": "tls_certs",
"description": "certificate renewal",
"content": "acme renews every 60 days; caddy reloads the SSLContext in place",
},
{
"name": "api_ratelimits",
"description": "rate limit policy",
"content": "100 req/min per token, burst 20, 429 with retry-after header",
},
{
"name": "backup_schedule",
"description": "nightly backups",
"content": "pg_dump at 02:00 UTC, retained 14 days, offsite copy weekly",
},
{
"name": "feature_flags",
"description": "rollout toggles",
"content": "rerank_bm25 default on, voice_io default off, smart_approvals off",
},
{
"name": "smtp_settings",
"description": "outbound email",
"content": "relay smtp.internal port 587 starttls from noreply@example.com",
},
{
"name": "log_retention",
"description": "log storage policy",
"content": "structured logs to loki, 30 day retention, audit logs 1 year",
},
{
"name": "node_placement",
"description": "routing config",
"content": "rendezvous hashing fnv-1a, ~100 node ceiling, hrw weights",
},
{
"name": "jwt_secret_rotation",
"description": "auth secret",
"content": "TURNSTONE_JWT_SECRET in config.toml, rotate quarterly, hs256",
},
]
MEMORY_LABELS: list[tuple[str, str]] = [
("what is the postgres database host and port", "postgres_conn"),
("how often do we rotate the jwt signing secret", "jwt_secret_rotation"),
("when do nightly database backups run", "backup_schedule"),
("who do I escalate an incident to", "oncall_rotation"),
("how are nodes placed for routing", "node_placement"),
("what is the per token api rate limit", "api_ratelimits"),
]
def _doc_text(d: dict[str, str]) -> str:
return " ".join(
filter(None, (d.get("name", ""), d.get("description", ""), d.get("content", "")))
)
def _make_rank(client: RerankClient, threshold: float) -> Callable[[str, list[str]], list[int]]:
def _rank(query: str, docs: list[str]) -> list[int]:
return [
h.index for h in client.rerank(query, docs) if threshold <= 0 or h.score >= threshold
]
return _rank
def _precision_at_k(result_names: list[str], relevant: str, k: int) -> float:
return 1.0 / k if relevant in result_names[:k] else 0.0
def _reciprocal_rank(result_names: list[str], relevant: str) -> float:
for i, name in enumerate(result_names, 1):
if name == relevant:
return 1.0 / i
return 0.0
def _bench_corpus(
title: str,
docs: list[dict[str, str]],
labels: list[tuple[str, str]],
client: RerankClient,
threshold: float,
k: int,
) -> None:
print(f"\n=== {title} ({len(docs)} docs, {len(labels)} queries, k={k}) ===")
texts = [_doc_text(d) for d in docs]
names = [d["name"] for d in docs]
plain = BM25Index(texts)
# Filter mode when a floor is set (mirrors the memory composition call site)
# so --threshold actually suppresses below-floor hits rather than being
# masked by reorder-mode backfill; reorder mode at threshold 0.
reranked = BM25Index(
texts, reranker=_make_rank(client, threshold), rerank_filters=threshold > 0
)
bm25_p = bm25_mrr = rr_p = rr_mrr = 0.0
for query, relevant in labels:
b_names = [names[i] for i in plain.search(query, k=k)]
r_names = [names[i] for i in reranked.search(query, k=k)]
bm25_p += _precision_at_k(b_names, relevant, k)
rr_p += _precision_at_k(r_names, relevant, k)
bm25_mrr += _reciprocal_rank([names[i] for i in plain.search(query, k=len(docs))], relevant)
rr_mrr += _reciprocal_rank(
[names[i] for i in reranked.search(query, k=len(docs))], relevant
)
flag = "" if b_names[:k] == r_names[:k] else " <-- reordered"
print(f" q: {query!r}")
print(f" want={relevant} bm25={b_names[:k]} rerank={r_names[:k]}{flag}")
n = len(labels)
print(f" -- precision@{k}: bm25={bm25_p / n:.3f} rerank={rr_p / n:.3f}")
print(f" -- MRR: bm25={bm25_mrr / n:.3f} rerank={rr_mrr / n:.3f}")
def _score_distribution(
docs: list[dict[str, str]],
labels: list[tuple[str, str]],
client: RerankClient,
) -> None:
"""Print rerank-score stats for labeled relevant vs irrelevant pairs.
A threshold default should sit above the irrelevant max / below the relevant
min where those separate; this prints both so you can eyeball the gap.
"""
texts = [_doc_text(d) for d in docs]
names = [d["name"] for d in docs]
relevant_scores: list[float] = []
irrelevant_scores: list[float] = []
for query, relevant in labels:
for hit in client.rerank(query, texts):
bucket = relevant_scores if names[hit.index] == relevant else irrelevant_scores
bucket.append(hit.score)
print("\n=== rerank score distribution (memory corpus) ===")
for label, scores in (("relevant", relevant_scores), ("irrelevant", irrelevant_scores)):
if not scores:
print(f" {label}: (no scores)")
continue
print(
f" {label:<10} n={len(scores):>3} "
f"min={min(scores):.4f} median={statistics.median(scores):.4f} "
f"max={max(scores):.4f}"
)
if relevant_scores and irrelevant_scores:
suggested = (min(relevant_scores) + max(irrelevant_scores)) / 2
sep = "separable" if min(relevant_scores) > max(irrelevant_scores) else "overlapping"
print(f" -> classes are {sep}; midpoint threshold candidate ~= {suggested:.4f}")
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--rerank-url", required=True, help="full Cohere/Jina-compatible /rerank URL")
ap.add_argument("--rerank-model", default="", help="model name sent in the request body")
ap.add_argument("--threshold", type=float, default=0.0, help="relevance floor (0 disables)")
ap.add_argument("--k", type=int, default=3, help="top-k to score precision@k over")
args = ap.parse_args()
# Bearer token comes from the environment, not a flag, to keep it out of
# shell history and the process listing.
api_key = os.environ.get("TURNSTONE_RERANK_API_KEY", "")
client = resolve_rerank_client(args.rerank_url, model=args.rerank_model, api_key=api_key)
if client is None:
print("No rerank endpoint resolved (empty --rerank-url?). Nothing to do.")
return 1
print(
f"reranker: url={args.rerank_url} model={args.rerank_model or '(default)'} "
f"threshold={args.threshold}"
)
_bench_corpus("tool search", TOOL_DOCS, TOOL_LABELS, client, args.threshold, args.k)
_bench_corpus("memory composition", MEMORY_DOCS, MEMORY_LABELS, client, args.threshold, args.k)
_score_distribution(MEMORY_DOCS, MEMORY_LABELS, client)
return 0
if __name__ == "__main__":
raise SystemExit(main())
+177 -1
View File
@@ -1,6 +1,6 @@
"""Tests for turnstone.core.bm25 — tokenizer and BM25 index."""
from turnstone.core.bm25 import BM25Index, _tokenize
from turnstone.core.bm25 import _RERANK_POOL, BM25Index, _tokenize
class TestTokenize:
@@ -69,3 +69,179 @@ class TestBM25Index:
results = index.search("file operations", k=3)
# Doc 2 has more file/operations mentions, should rank higher
assert results[0] == 2
class TestBM25Reranking:
"""Two-stage search: BM25 recall pool reordered by an attached reranker.
The reranker is a plain callable returning POSITIONS into the recall pool
(best-first); ``search`` maps each position back to the original doc index
via ``pool[pos]``. Fakes are deterministic lambdas/closures — the rerank
HTTP boundary is exercised separately in test_rerank.py.
"""
# All matching docs share "alpha"; the non-matching ones (zzz) stay out of
# the pool, so a returned value that is 1 or 3 would prove a mapping bug.
_DOCS = ["alpha one", "zzz", "alpha two", "zzz", "alpha three"]
def test_reorders_by_reranker_output(self):
# Reranker dictates pool order 2, 0, 1 (positions) regardless of BM25.
docs = ["alpha", "alpha", "alpha"]
index = BM25Index(docs, reranker=lambda q, d: [2, 0, 1])
# pool == [0, 1, 2] (BM25 tie-break is ascending index); positions map
# back to original indices 2, 0, 1.
assert index.search("alpha", k=3) == [2, 0, 1]
def test_maps_reranker_position_back_to_original_index(self):
# Reranker reverses the pool it is handed (positions n-1..0). Guards
# the ``pool[pos]`` mapping line: returned values must be ORIGINAL doc
# indices (a subset of {0, 2, 4}), never raw positions like 1 or 3.
index = BM25Index(self._DOCS, reranker=lambda q, d: list(range(len(d)))[::-1])
result = index.search("alpha", k=5)
# Only the three "alpha" docs are in the pool.
assert set(result) == {0, 2, 4}
# Reversed pool ordering: whatever BM25 pool order was, it is reversed.
bm25_pool = BM25Index(self._DOCS).search("alpha", k=_RERANK_POOL)
assert result == bm25_pool[::-1]
def test_exception_falls_back_to_bm25_order(self):
def boom(q, d):
raise RuntimeError("rerank endpoint down")
index = BM25Index(self._DOCS, reranker=boom)
# Guards the ``except`` clause -> BM25 order, top-k.
assert index.search("alpha", k=2) == BM25Index(self._DOCS)._bm25_rank("alpha")[:2]
def test_clean_empty_return_is_honored_no_fallback(self):
# FILTER MODE (rerank_filters=True, the memory floor): a clean empty
# return means "inject nothing" (the caller's relevance floor emptied
# it) and must NOT fall back to BM25 order. Guards the filter-branch
# ``return out`` -- adding ``if not out: return pool[:k]`` there fails.
index = BM25Index(self._DOCS, reranker=lambda q, d: [], rerank_filters=True)
assert index.search("alpha", k=5) == []
def test_filter_mode_exception_still_falls_back(self):
# FILTER MODE honors a clean empty (above), but a reranker EXCEPTION is
# an endpoint failure, not a floor verdict -> BM25 fallback, NOT empty.
# This is the parse-failure-vs-floor distinction at the seam: the
# _bm25_reranker closure raises on an unparseable/empty response so
# memory composition can't be silently suppressed by a broken endpoint.
def boom(q, d):
raise RuntimeError("rerank endpoint down")
index = BM25Index(self._DOCS, reranker=boom, rerank_filters=True)
assert index.search("alpha", k=2) == BM25Index(self._DOCS)._bm25_rank("alpha")[:2]
def test_singleton_pool_still_reranked(self):
# A 1-doc pool is still sent to the reranker (no len>1 short-circuit).
# Recording reranker returns [] -> result is [] AND it WAS called once.
# Filter mode honors the empty result (no fallback).
calls = {"n": 0}
def rec(q, d):
calls["n"] += 1
return []
# "two" matches only docs[2] -> pool of size 1.
index = BM25Index(self._DOCS, reranker=rec, rerank_filters=True)
assert index.search("two", k=5) == []
assert calls["n"] == 1
def test_reorder_mode_empty_falls_back_to_bm25(self):
# REORDER MODE (default rerank_filters=False, reactive tool/skill
# search): an empty reranker result means the endpoint failed, so fall
# back to BM25 order -- results are NEVER silently dropped. [bug-1]
# Guards the reorder-branch ``if not out: return pool[:k]`` -- removing
# it makes this return [] and fail.
index = BM25Index(self._DOCS, reranker=lambda q, d: [])
assert index.search("alpha", k=5) == BM25Index(self._DOCS)._bm25_rank("alpha")[:5]
def test_reorder_mode_backfills_omitted_pool_items(self):
# REORDER MODE: the reranker returns a STRICT SUBSET of pool positions
# (only the first two). The two reranked items come FIRST, then the
# omitted pool items are backfilled in BM25 order, capped at k. [bug-2]
# Guards the backfill loop -- removing it drops the omitted docs.
bm25_pool = BM25Index(self._DOCS)._bm25_rank("alpha") # 3 matching docs
assert len(bm25_pool) == 3
# Reranker keeps only pool positions 0 and 1 (drops the third).
index = BM25Index(self._DOCS, reranker=lambda q, d: [0, 1])
result = index.search("alpha", k=5)
# Reranked two first (pool[0], pool[1]) ...
assert result[:2] == [bm25_pool[0], bm25_pool[1]]
# ... then the omitted pool item backfilled in BM25 order.
assert result == [bm25_pool[0], bm25_pool[1], bm25_pool[2]]
# No item silently lost: the full matching set is present.
assert set(result) == set(bm25_pool)
def test_reorder_mode_backfill_respects_k(self):
# Backfill must stop at k: subset rerank + a k smaller than the pool.
bm25_pool = BM25Index(self._DOCS)._bm25_rank("alpha") # 3 matching docs
index = BM25Index(self._DOCS, reranker=lambda q, d: [0])
result = index.search("alpha", k=2)
# One reranked item, then one backfilled, capped at k=2.
assert result == [bm25_pool[0], bm25_pool[1]]
def test_reorder_mode_full_reorder_unchanged(self):
# REORDER MODE (default) with a full permutation: behaves exactly like
# the historical reorder -- every pool item present, in reranker order,
# backfill loop adds nothing (all positions already seen).
index = BM25Index(self._DOCS, reranker=lambda q, d: list(range(len(d)))[::-1])
result = index.search("alpha", k=5)
bm25_pool = BM25Index(self._DOCS).search("alpha", k=_RERANK_POOL)
assert result == bm25_pool[::-1]
def test_no_reranker_matches_baseline(self):
docs = [
"read a file from disk",
"search for file in directory",
"execute a bash command about files",
"totally unrelated cooking content",
]
plain = BM25Index(docs)
attached = BM25Index(docs, reranker=lambda q, d: list(range(len(d))))
for q in ("file", "bash command", "disk", "cooking", "file directory"):
# reranker=identity returns the pool unchanged, but the no-reranker
# path must be byte-for-byte the historical result regardless.
assert plain.search(q, k=3) == BM25Index(docs).search(q, k=3)
# And identity-rerank reproduces the BM25 top-k for these queries.
assert attached.search(q, k=3) == plain.search(q, k=3)
def test_bool_positions_rejected(self):
docs = ["alpha", "alpha", "alpha"]
# True/False are int subclasses posing as 1/0 -> rejected; "1" is not an
# int -> rejected; only positions 2 and 0 survive. Filter mode so the
# rejected positions are not re-added by reorder-mode backfill -- this
# isolates the type guard.
index = BM25Index(docs, reranker=lambda q, d: [True, "1", 2, 0], rerank_filters=True)
assert index.search("alpha", k=5) == [2, 0]
def test_recall_pool_capped_at_rerank_pool(self):
# More matching docs than the recall cap: the reranker must receive
# exactly _RERANK_POOL docs, and outputs may only reference that set.
n = _RERANK_POOL + 10
docs = [f"alpha doc{i}" for i in range(n)]
seen_len = {"n": -1}
def rec(q, d):
seen_len["n"] = len(d)
return list(range(len(d))) # identity over the (capped) pool
index = BM25Index(docs, reranker=rec)
result = index.search("alpha", k=n)
assert seen_len["n"] == _RERANK_POOL # only the first 50 reached rerank
assert len(result) == _RERANK_POOL
# Every returned index is from the BM25 top-50 recall set.
recall = set(BM25Index(docs)._bm25_rank("alpha")[:_RERANK_POOL])
assert set(result) == recall
def test_empty_query_skips_reranker(self):
calls = {"n": 0}
def rec(q, d):
calls["n"] += 1
return list(range(len(d)))
index = BM25Index(self._DOCS, reranker=rec)
# Empty query -> empty BM25 pool -> reranker never invoked.
assert index.search("", k=5) == []
assert calls["n"] == 0
+108
View File
@@ -75,6 +75,73 @@ class TestScoreMemories:
assert result[0]["name"] == "x"
class TestScoreMemoriesReranking:
"""``score_memories`` forwards a reranker into the BM25 recall pool.
The reranker is a deterministic callable over POSITIONS in the recall pool
(the matched memories, BM25-ordered); the result is the corresponding memory
dicts, best-first. The existing 7 tests above pass no reranker (default
None) and exercise the unchanged BM25-only path.
"""
_MEMS = [
{"name": "alpha", "description": "shared topic", "content": "shared topic alpha"},
{"name": "beta", "description": "shared topic", "content": "shared topic beta"},
{"name": "gamma", "description": "shared topic", "content": "shared topic gamma"},
]
def test_reranker_reorders_memories(self):
# All three match "shared topic" -> pool covers them. The reranker
# reverses the pool positions, so the returned memory order is the
# BM25 order reversed.
baseline = score_memories(self._MEMS, "shared topic", k=3)
reranked = score_memories(
self._MEMS,
"shared topic",
k=3,
reranker=lambda q, d: list(range(len(d)))[::-1],
)
assert [m["name"] for m in reranked] == [m["name"] for m in baseline][::-1]
# Still the same set of memories, just reordered.
assert {m["name"] for m in reranked} == {m["name"] for m in baseline}
def test_floor_empties_returns_nothing(self):
# FILTER MODE (rerank_filters=True): a relevance floor that rejects
# everything (reranker returns []) means "inject no memory" ->
# score_memories returns []. This is the proactive memory floor the
# threshold setting drives (an active floor -> rerank_filters=True).
result = score_memories(
self._MEMS,
"shared topic",
k=3,
reranker=lambda q, d: [],
rerank_filters=True,
)
assert result == []
def test_reorder_mode_empty_does_not_suppress(self):
# REORDER MODE (rerank_filters=False, the disabled-floor default): an
# empty reranker result means the endpoint failed, NOT "suppress all".
# Memories fall back to BM25 top-k -- never silently dropped. Guards the
# threshold<=0 -> reorder-mode wiring in the memory call site.
result = score_memories(
self._MEMS,
"shared topic",
k=3,
reranker=lambda q, d: [],
rerank_filters=False,
)
baseline = score_memories(self._MEMS, "shared topic", k=3)
assert [m["name"] for m in result] == [m["name"] for m in baseline]
assert len(result) == 3
def test_default_none_unchanged(self):
# No reranker kwarg -> identical to passing reranker=None -> BM25-only.
assert score_memories(self._MEMS, "shared topic", k=2) == score_memories(
self._MEMS, "shared topic", k=2, reranker=None
)
# ---------------------------------------------------------------------------
# build_memory_context
# ---------------------------------------------------------------------------
@@ -395,6 +462,47 @@ class TestCompositionCandidateSelection:
assert search_mock.call_args.args[1] == [("coordinator", "coord-1")]
class TestCompositionRerankFiltersWiring:
"""The memory composition call site maps ``threshold > 0`` to ``rerank_filters``.
A disabled floor (threshold <= 0) -> reorder mode (rerank_filters=False) so an
empty/failed reranker falls back to BM25 (memories not suppressed); an active
floor (threshold > 0) -> filter mode (rerank_filters=True) so the floor may
legitimately empty the injection. Drives the real ``_init_system_messages``
call site, capturing the kwarg ``score_memories`` actually receives.
"""
def _capture_rerank_filters(self, session: object, threshold: float) -> bool:
captured: dict[str, bool] = {}
def _fake_score(*_args: object, rerank_filters: bool = False, **_kw: object):
captured["rerank_filters"] = rerank_filters
return []
mem = _make_mem("m_one", content="alpha")
with (
patch("turnstone.core.session.score_memories", _fake_score),
patch.object(session, "_bm25_rerank_threshold", return_value=threshold),
patch.object(session, "_bm25_reranker", return_value=None),
patch.object(session, "_select_memory_candidates", return_value=([mem], "list")),
):
session._init_system_messages()
assert "rerank_filters" in captured, "score_memories was not reached"
return captured["rerank_filters"]
def test_threshold_zero_uses_reorder_mode(self, tmp_db):
session = _make_session()
session.messages = [{"role": "user", "content": "alpha"}]
# threshold 0 (disabled floor) -> reorder mode -> no suppression.
assert self._capture_rerank_filters(session, 0.0) is False
def test_positive_threshold_uses_filter_mode(self, tmp_db):
session = _make_session()
session.messages = [{"role": "user", "content": "alpha"}]
# An active floor -> filter mode -> the reranker may empty the injection.
assert self._capture_rerank_filters(session, 0.5) is True
class TestMemorySearchToolExecution:
"""End-to-end test of ``memory(action='search')`` through _exec_memory.
+112
View File
@@ -301,3 +301,115 @@ class TestSessionRerankWiring:
_registry=SimpleNamespace(get_config=lambda a: cfg),
)
assert ChatSession._resolve_rerank_client(stub) is None # falls through, no rerank_url
# ---------------------------------------------------------------------------
# ChatSession BM25 reranker — the closure feeding tool/skill/memory retrieval
# ---------------------------------------------------------------------------
class _FakeRerankClient:
"""In-process RerankClient stub returning fixed hits (no HTTP)."""
def __init__(self, hits: list[RerankHit]) -> None:
self._hits = hits
def rerank(
self, query: str, documents: list[str], *, top_n: int | None = None
) -> list[RerankHit]:
return self._hits
class TestSessionBM25Reranker:
"""``_bm25_reranker`` / ``_bm25_rerank_threshold`` — the BM25 seam adapters.
Drives the real closure through a fake ``RerankClient`` (the in-process
callable seam), never by patching internal state. The HTTP boundary lives in
``TestCohereJinaRerankClient`` above.
"""
def test_none_when_disabled_for_bm25(self):
# Per-tool toggle off -> no reranker even with an endpoint configured.
stub = SimpleNamespace(
_rerank_enabled_for=lambda tool: False,
_resolve_rerank_client=lambda: _FakeRerankClient([]),
)
assert ChatSession._bm25_reranker(stub) is None
def test_none_when_no_client(self):
# Enabled, but no endpoint resolves -> None.
stub = SimpleNamespace(
_rerank_enabled_for=lambda tool: True,
_resolve_rerank_client=lambda: None,
)
assert ChatSession._bm25_reranker(stub) is None
def _enabled_stub(self, hits: list[RerankHit]) -> SimpleNamespace:
return SimpleNamespace(
_rerank_enabled_for=lambda tool: True,
_resolve_rerank_client=lambda: _FakeRerankClient(hits),
)
def test_no_threshold_returns_all_hit_indices(self):
hits = [RerankHit(index=0, score=0.9), RerankHit(index=1, score=0.1)]
rank = ChatSession._bm25_reranker(self._enabled_stub(hits), 0.0)
assert rank is not None
# threshold 0 disables the floor -> every hit index passes through.
assert rank("q", ["d0", "d1"]) == [0, 1]
def test_threshold_filters_below_floor(self):
hits = [RerankHit(index=0, score=0.9), RerankHit(index=1, score=0.1)]
rank = ChatSession._bm25_reranker(self._enabled_stub(hits), 0.5)
assert rank is not None
# Only the 0.9 hit clears the 0.5 floor. Guards the ``h.score >=
# threshold`` boundary (flip to ``>`` / ``<`` and idx1 leaks or idx0
# drops).
assert rank("q", ["d0", "d1"]) == [0]
def test_threshold_boundary_is_inclusive(self):
# A hit exactly at the floor is KEPT (>= , not >).
hits = [RerankHit(index=0, score=0.5)]
rank = ChatSession._bm25_reranker(self._enabled_stub(hits), 0.5)
assert rank is not None
assert rank("q", ["d0"]) == [0]
def test_empty_hits_raise_for_nonempty_docs(self):
# A conforming reranker scores every doc; [] for non-empty input is an
# endpoint/parse failure, NOT a floor result -> raise (a discrete branch
# from the threshold) so BM25Index falls back to BM25 order in BOTH
# modes. Holds regardless of threshold.
from turnstone.core.rerank import RerankError
for thr in (0.0, 0.5):
rank = ChatSession._bm25_reranker(self._enabled_stub([]), thr)
assert rank is not None
with pytest.raises(RerankError):
rank("q", ["d0", "d1"])
def test_floor_dropping_all_returns_empty_not_raise(self):
# Distinct from a parse failure: the reranker DID score the doc, the
# floor just dropped it -> clean empty (honored by filter mode), no raise.
hits = [RerankHit(index=0, score=0.1)]
rank = ChatSession._bm25_reranker(self._enabled_stub(hits), 0.5)
assert rank is not None
assert rank("q", ["d0"]) == []
def test_empty_docs_does_not_raise(self):
# Empty input legitimately yields empty output -- nothing to signal.
rank = ChatSession._bm25_reranker(self._enabled_stub([]), 0.0)
assert rank is not None
assert rank("q", []) == []
def test_threshold_reads_setting(self):
cs = _FakeConfigStore({"tools.rerank_bm25_threshold": 0.42})
stub = SimpleNamespace(_config_store=cs)
assert ChatSession._bm25_rerank_threshold(stub) == 0.42
def test_threshold_zero_without_config_store(self):
stub = SimpleNamespace(_config_store=None)
assert ChatSession._bm25_rerank_threshold(stub) == 0.0
def test_threshold_zero_on_garbage_value(self):
cs = _FakeConfigStore({"tools.rerank_bm25_threshold": "not-a-number"})
stub = SimpleNamespace(_config_store=cs)
assert ChatSession._bm25_rerank_threshold(stub) == 0.0
+4
View File
@@ -75,6 +75,10 @@ def _make_session(*, kind: str = "interactive", user_id: str = "test-user") -> A
)
# Truncation budget — required by _truncate_output on every exec.
session.tool_truncation = 100_000
# skills(action='find') ranks via BM25Index(..., reranker=self._bm25_reranker()),
# which reaches _resolve_rerank_client -> self.tool_timeout. No _config_store/
# _registry here -> no endpoint -> reranker is None -> pure-BM25 path.
session.tool_timeout = 30
# set_skill stub for load action. Records both the name and the
# arguments-string so tests can assert that #572's invocation-args
+33
View File
@@ -218,6 +218,39 @@ class TestToolSearchManager:
assert "mcp__github__create_issue" in text
class TestToolSearchManagerReranking:
"""``ToolSearchManager`` forwards a reranker into its deferred-tool index."""
def _tools(self):
# All three deferred tools match "github" so the recall pool spans them;
# a reranker can then dictate their order.
return [
_make_tool("bash", "Execute shell commands"),
_make_tool("mcp__github__a", "github alpha helper"),
_make_tool("mcp__github__b", "github beta helper"),
_make_tool("mcp__github__c", "github gamma helper"),
]
def test_search_reflects_reranker_order(self):
baseline = ToolSearchManager(self._tools(), always_on_names={"bash"})
base_names = [_tool_name(t) for t in baseline.search("github helper")]
# Reranker reverses the recall-pool order it is handed (positions
# n-1..0). The forwarded order must show up in search results.
reranked = ToolSearchManager(
self._tools(),
always_on_names={"bash"},
reranker=lambda q, d: list(range(len(d)))[::-1],
)
names = [_tool_name(t) for t in reranked.search("github helper")]
assert names == base_names[::-1]
def test_no_reranker_unchanged(self):
mgr = ToolSearchManager(self._tools(), always_on_names={"bash"}, reranker=None)
names = {_tool_name(t) for t in mgr.search("github helper")}
assert names == {"mcp__github__a", "mcp__github__b", "mcp__github__c"}
# ---------------------------------------------------------------------------
# Helper function tests
# ---------------------------------------------------------------------------
+1 -1
View File
@@ -5205,7 +5205,7 @@ const MODEL_ROLES = [
{
label: "Reranker",
description:
"Reranks web_search results. Point at a model whose base_url is a Cohere/Jina-compatible /rerank endpoint and whose capabilities include supports_rerank. Empty falls back to the tools.rerank_url setting (if set).",
"Reranks web_search results. Point at a model whose base_url is a Cohere/Jina-compatible /rerank endpoint and whose capabilities include supports_rerank. Empty falls back to the tools.rerank_url setting (if set). Enabling a reranker sends web_search results AND BM25 retrieval candidates (tool/skill descriptions and memory content) to this endpoint; self-hosted endpoints keep it on your infrastructure.",
aliasKey: "tools.reranker_alias",
fallbackKind: "disabled",
mediaCapability: "supports_rerank",
+75 -4
View File
@@ -8,9 +8,19 @@ from __future__ import annotations
import math
import re
from collections import Counter
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from turnstone.core.rerank import Reranker
_SPLIT_RE = re.compile(r"[_\-./\s]+")
# Max BM25 hits handed to the reranker in one request — the recall pool whose
# top-k the reranker reorders before the caller's ``k`` slice. Deliberately a
# private copy (not shared from rerank.py) so this module stays import-light and
# httpx-free; web_search.py defines the same cap independently — keep them in sync.
_RERANK_POOL = 50
def _tokenize(text: str) -> list[str]:
"""Split text on whitespace, underscores, hyphens, dots."""
@@ -20,9 +30,19 @@ def _tokenize(text: str) -> list[str]:
class BM25Index:
"""Okapi BM25 ranking index over short text documents."""
def __init__(self, documents: list[str], *, k1: float = 1.5, b: float = 0.75) -> None:
def __init__(
self,
documents: list[str],
*,
k1: float = 1.5,
b: float = 0.75,
reranker: Reranker | None = None,
rerank_filters: bool = False,
) -> None:
self.k1 = k1
self.b = b
self._reranker = reranker
self._rerank_filters = rerank_filters
self._docs = documents
self._doc_tokens: list[list[str]] = [_tokenize(d) for d in documents]
self._doc_lens = [len(t) for t in self._doc_tokens]
@@ -34,8 +54,8 @@ class BM25Index:
for term in set(tokens):
self._df[term] += 1
def search(self, query: str, k: int = 5) -> list[int]:
"""Return indices of top-k documents sorted by descending BM25 score."""
def _bm25_rank(self, query: str) -> list[int]:
"""Return ALL matching document indices sorted by descending BM25 score."""
q_tokens = _tokenize(query)
if not q_tokens:
return []
@@ -45,7 +65,58 @@ class BM25Index:
if score > 0:
scores.append((score, idx))
scores.sort(key=lambda x: (-x[0], x[1]))
return [idx for _, idx in scores[:k]]
return [idx for _, idx in scores]
def search(self, query: str, k: int = 5) -> list[int]:
"""Return indices of top-k documents by relevance.
Stage 1 is BM25. When a reranker is attached, a recall pool of the top
``_RERANK_POOL`` BM25 hits is reranked and the top-k spliced back.
"""
ranked = self._bm25_rank(query)
if self._reranker is None:
return ranked[:k] # no reranker: today's behavior, byte-for-byte
pool = ranked[:_RERANK_POOL]
if not pool:
return []
docs = [self._docs[i] for i in pool]
try:
order = list(self._reranker(query, docs))
except Exception:
return pool[:k] # reranker ERROR -> BM25 order (both modes)
seen: set[int] = set()
out: list[int] = []
for pos in order:
# bool is an int subclass -- reject a stray True/False posing as 1/0.
if (
isinstance(pos, int)
and not isinstance(pos, bool)
and 0 <= pos < len(pool)
and pos not in seen
):
seen.add(pos)
out.append(pool[pos]) # map reranker position -> original doc index
if len(out) >= k:
break
if self._rerank_filters:
# FILTER MODE (memory floor): the reranker may legitimately drop
# sub-floor candidates, so a clean short/empty result is HONORED as-is
# -- NO fallback to BM25 order, NO backfill. This is the deliberate
# divergence from web_search._rerank_results and from reorder mode below;
# do not "fix" it to fall back on empty (a test pins this).
return out
# REORDER MODE (reactive tool/skill search): the reranker must never drop
# candidates. An empty result means the endpoint failed -> BM25 fallback;
# any pool items the reranker omitted (e.g. a top_n subset) are backfilled
# in BM25 order so results are never silently lost.
if not out:
return pool[:k]
for pos in range(len(pool)):
if len(out) >= k:
break
if pos not in seen:
out.append(pool[pos])
return out
def _score(self, q_tokens: list[str], doc_tokens: list[str], dl: int) -> float:
tf_map: Counter[str] = Counter(doc_tokens)
+7 -2
View File
@@ -4,10 +4,13 @@ from __future__ import annotations
from dataclasses import dataclass
from html import escape as _html_escape
from typing import Any
from typing import TYPE_CHECKING, Any
from turnstone.core.bm25 import BM25Index
if TYPE_CHECKING:
from turnstone.core.rerank import Reranker
@dataclass
class MemoryConfig:
@@ -24,6 +27,8 @@ def score_memories(
memories: list[dict[str, str]],
query: str,
k: int = 5,
reranker: Reranker | None = None,
rerank_filters: bool = False,
) -> list[dict[str, str]]:
"""Return the top-k memories most relevant to *query*.
@@ -41,7 +46,7 @@ def score_memories(
f"{m.get('name', '')} {m.get('description', '')} {m.get('content', '')[:200]}"
for m in memories
]
index = BM25Index(documents)
index = BM25Index(documents, reranker=reranker, rerank_filters=rerank_filters)
top_indices = index.search(query, k)
return [memories[i] for i in top_indices]
+17
View File
@@ -21,6 +21,7 @@ default and no fall back to a local model.
from __future__ import annotations
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any, Protocol
@@ -30,6 +31,22 @@ from turnstone.core.log import get_logger
log = get_logger(__name__)
# A reranker reorders candidate documents by relevance to the query, returning
# their indices best-first. Defined here so bm25.py and web_search.py can share
# the type without importing each other (rerank.py imports neither).
Reranker = Callable[[str, list[str]], list[int]]
class RerankError(RuntimeError):
"""A rerank endpoint returned no usable scores for a non-empty input.
A conforming reranker scores every document, so an empty result for
non-empty input means the response was unparseable / non-conforming -- an
endpoint failure, distinct from a relevance floor dropping every candidate.
Callers raise this so retrieval falls back to BM25 order rather than
treating the failure as "nothing relevant".
"""
@dataclass(frozen=True)
class RerankHit:
+65 -7
View File
@@ -142,9 +142,9 @@ if TYPE_CHECKING:
ModelCapabilities,
StreamChunk,
)
from turnstone.core.rerank import RerankClient
from turnstone.core.rerank import RerankClient, Reranker
from turnstone.core.tool_advisory import ToolAdvisory
from turnstone.core.web_search import Reranker, WebSearchClient
from turnstone.core.web_search import WebSearchClient
# ---------------------------------------------------------------------------
# Cancellation support
@@ -604,6 +604,10 @@ _WATCH_QUEUE_SOFT_CAP = 50
# the codepoint count.
REMINDER_TEXT_STORAGE_CAP = 8192
_RERANK_TIMEOUT_CAP_S = 15.0 # reranking <=50 short docs is fast; cap so a hung
# endpoint falls back to BM25 in seconds, not up to tools.timeout (120s default).
# Per-turn memory rerank makes the long timeout a turn-stall hazard.
def _without_tool(tools: list[dict[str, Any]], name: str) -> list[dict[str, Any]]:
"""Return *tools* with the named tool removed."""
@@ -1178,6 +1182,7 @@ class ChatSession:
self._tools,
always_on_names=builtin_in_session,
max_results=tool_search_max_results,
reranker=self._bm25_reranker(),
)
# Skill: explicit name overrides is_default skills. ``skill_arguments``
# carries the spec's $ARGUMENTS payload — set at create/load time,
@@ -1339,7 +1344,7 @@ class ChatSession:
url=cfg.base_url,
model=cfg.model or "",
api_key=cfg.api_key,
timeout=self.tool_timeout,
timeout=min(self.tool_timeout, _RERANK_TIMEOUT_CAP_S),
)
def _setting(key: str, env_value: str) -> str:
@@ -1355,11 +1360,11 @@ class ChatSession:
url=_setting("tools.rerank_url", get_rerank_url()),
model=_setting("tools.rerank_model", get_rerank_model()),
api_key=_setting("tools.rerank_api_key", get_rerank_api_key()),
timeout=self.tool_timeout,
timeout=min(self.tool_timeout, _RERANK_TIMEOUT_CAP_S),
)
def _rerank_enabled_for(self, tool: str) -> bool:
"""Whether reranking is enabled for ``tool`` (currently: 'web_search').
"""Whether reranking is enabled for ``tool`` (currently: 'web_search', 'bm25').
The per-tool toggles default on; the operative gate is whether an
endpoint is configured (``_resolve_rerank_client`` returns None when
@@ -1387,6 +1392,51 @@ class ChatSession:
return _rank
def _bm25_reranker(self, threshold: float = 0.0) -> Reranker | None:
"""Build a BM25 reranker callable, or None when disabled.
Mirrors ``_web_search_reranker``. ``threshold`` is a relevance FLOOR
applied in this closure (where scores still exist); the BM25Index seam
stays indices-only. ``threshold <= 0`` disables the floor. An empty
response for non-empty input is an endpoint failure (a conforming
reranker scores every doc), NOT a floor result, so it raises
``RerankError`` -> BM25Index falls back to BM25 order regardless of
threshold. Only memory composition passes a configured threshold;
reactive surfaces pass 0.
"""
if not self._rerank_enabled_for("bm25"):
return None
rc = self._resolve_rerank_client()
if rc is None:
return None
from turnstone.core.rerank import RerankError
def _rank(query: str, docs: list[str]) -> list[int]:
hits = rc.rerank(query, docs)
if docs and not hits:
# A conforming reranker scores every document; an empty result
# for non-empty input means the endpoint response was
# unparseable (rc.rerank -> _parse_hits returns [] without
# raising). That is an endpoint FAILURE -- a discrete branch
# from the relevance floor below -- so raise and let BM25Index
# fall back to BM25 order in BOTH modes, instead of the
# filter-mode floor honoring it as "nothing relevant".
raise RerankError("rerank endpoint returned no scores for non-empty input")
return [h.index for h in hits if threshold <= 0 or h.score >= threshold]
return _rank
def _bm25_rerank_threshold(self) -> float:
"""Configured proactive-memory relevance floor (0.0 = disabled)."""
cs = getattr(self, "_config_store", None)
if cs is None:
return 0.0
try:
return float(cs.get("tools.rerank_bm25_threshold") or 0.0)
except (TypeError, ValueError):
return 0.0
def _resolve_capabilities(
self,
provider: LLMProvider,
@@ -2064,6 +2114,7 @@ class ChatSession:
self._tools,
always_on_names=set(BUILTIN_TOOL_NAMES),
max_results=self._tool_search_max_results,
reranker=self._bm25_reranker(),
)
# Restore previously expanded tools that still exist
if old_expanded:
@@ -2767,7 +2818,14 @@ class ChatSession:
context = extract_recent_context(self.messages)
visible_mems, candidate_source = self._select_memory_candidates(context)
if visible_mems:
relevant = score_memories(visible_mems, context, k=self._mem_cfg.relevance_k)
thr = self._bm25_rerank_threshold()
relevant = score_memories(
visible_mems,
context,
k=self._mem_cfg.relevance_k,
reranker=self._bm25_reranker(thr),
rerank_filters=thr > 0,
)
log.info(
"memory.composition",
source=candidate_source,
@@ -8735,7 +8793,7 @@ class ChatSession:
)
for r in rows
]
index = BM25Index(corpus)
index = BM25Index(corpus, reranker=self._bm25_reranker())
top = index.search(query, k=min(len(rows), 50))
rows = [rows[i] for i in top]
skills = [self._skills_project_row(r) for r in rows]
+24
View File
@@ -307,6 +307,30 @@ def _build_registry() -> dict[str, SettingDef]:
"compatible /rerank endpoint, then pick it under Models -> Roles -> Reranker. Takes "
"precedence over tools.rerank_url; empty falls back to tools.rerank_url (if set).",
),
SettingDef(
"tools.rerank_bm25",
"bool",
True,
"Rerank BM25-backed retrieval (tool/skill search, memory) when configured",
"tools",
help="When a rerank endpoint is configured, rerank BM25-backed retrieval "
"(tool search, skill search, memory composition). Disable on low-power hosts "
"to keep web_search reranking without paying the per-turn memory-composition "
"rerank. Reranking sends the candidate text (tool/skill names + descriptions, "
"and memory name/description/content) to the configured rerank endpoint; "
"self-hosted (vLLM/llama.cpp/TEI) keeps it on your infrastructure, a hosted "
"provider (Cohere/Jina/Voyage) sends it off-box.",
),
SettingDef(
"tools.rerank_bm25_threshold",
"float",
0.0,
"Relevance floor (0-1) for proactive memory surfacing; 0 disables",
"tools",
help="0-1 relevance floor for PROACTIVE memory surfacing; 0 disables it. On "
"your reranker's score scale -- calibrated 0-1 on Cohere/Jina/Qwen endpoints, "
"raw logits on bge/TEI.",
),
# -- server ---------------------------------------------------------
SettingDef(
"server.workstream_idle_timeout",
+6 -2
View File
@@ -10,10 +10,13 @@ from __future__ import annotations
import re
from collections import Counter
from typing import Any
from typing import TYPE_CHECKING, Any
from turnstone.core.bm25 import BM25Index, _tokenize # noqa: F401
if TYPE_CHECKING:
from turnstone.core.rerank import Reranker
# ---------------------------------------------------------------------------
# Tool search manager — partitions tools, tracks visibility
# ---------------------------------------------------------------------------
@@ -65,6 +68,7 @@ class ToolSearchManager:
always_on_names: set[str],
*,
max_results: int = 5,
reranker: Reranker | None = None,
) -> None:
self._always_on: list[dict[str, Any]] = []
self._deferred: list[dict[str, Any]] = []
@@ -82,7 +86,7 @@ class ToolSearchManager:
# BM25 index over deferred tools
texts = [_tool_text(t) for t in self._deferred]
self._index = BM25Index(texts)
self._index = BM25Index(texts, reranker=reranker)
# Pre-compute server summary for the search tool description
self._server_hint = _mcp_server_summary(self._deferred)
+3 -7
View File
@@ -16,7 +16,6 @@ clients.
from __future__ import annotations
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, Protocol
import httpx
@@ -25,16 +24,13 @@ from turnstone.core.log import get_logger
if TYPE_CHECKING:
from turnstone.core.mcp_client import MCPClientManager
from turnstone.core.rerank import Reranker
log = get_logger(__name__)
# A reranker reorders candidate documents by relevance to the query, returning
# their indices best-first. ``web_search`` accepts one so the session can plug
# in an endpoint-backed reranker without this module depending on rerank.py.
Reranker = Callable[[str, list[str]], list[int]]
# Max SearxNG results sent to the reranker in one request — the pool re-ordered
# before the caller's ``max_results`` slice.
# before the caller's ``max_results`` slice. bm25.py defines the same cap
# independently (kept separate so bm25 stays httpx-free) — keep the two in sync.
_RERANK_POOL = 50