From 215f7506badb3d82f9e41a6ecde2e6af54cbb865 Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Mon, 1 Jun 2026 12:40:39 -0700 Subject: [PATCH] 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. --- docs/tools.md | 2 + scripts/bench_bm25_rerank.py | 260 ++++++++++++++++++++++++++++ tests/test_bm25.py | 178 ++++++++++++++++++- tests/test_memory_relevance.py | 108 ++++++++++++ tests/test_rerank.py | 112 ++++++++++++ tests/test_skills_tool.py | 4 + tests/test_tool_search.py | 33 ++++ turnstone/console/static/admin.js | 2 +- turnstone/core/bm25.py | 79 ++++++++- turnstone/core/memory_relevance.py | 9 +- turnstone/core/rerank.py | 17 ++ turnstone/core/session.py | 72 +++++++- turnstone/core/settings_registry.py | 24 +++ turnstone/core/tool_search.py | 8 +- turnstone/core/web_search.py | 10 +- 15 files changed, 894 insertions(+), 24 deletions(-) create mode 100644 scripts/bench_bm25_rerank.py diff --git a/docs/tools.md b/docs/tools.md index b0ed131a..0a392d1f 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -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 diff --git a/scripts/bench_bm25_rerank.py b/scripts/bench_bm25_rerank.py new file mode 100644 index 00000000..b96f7899 --- /dev/null +++ b/scripts/bench_bm25_rerank.py @@ -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()) diff --git a/tests/test_bm25.py b/tests/test_bm25.py index 2ee97595..625cf26a 100644 --- a/tests/test_bm25.py +++ b/tests/test_bm25.py @@ -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 diff --git a/tests/test_memory_relevance.py b/tests/test_memory_relevance.py index b5f043e4..4bfce98e 100644 --- a/tests/test_memory_relevance.py +++ b/tests/test_memory_relevance.py @@ -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. diff --git a/tests/test_rerank.py b/tests/test_rerank.py index 01615eb5..cc797490 100644 --- a/tests/test_rerank.py +++ b/tests/test_rerank.py @@ -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 diff --git a/tests/test_skills_tool.py b/tests/test_skills_tool.py index 3923e7c1..432de776 100644 --- a/tests/test_skills_tool.py +++ b/tests/test_skills_tool.py @@ -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 diff --git a/tests/test_tool_search.py b/tests/test_tool_search.py index e08e4039..5cfa9a7a 100644 --- a/tests/test_tool_search.py +++ b/tests/test_tool_search.py @@ -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 # --------------------------------------------------------------------------- diff --git a/turnstone/console/static/admin.js b/turnstone/console/static/admin.js index b17928a2..d9fcc7ad 100644 --- a/turnstone/console/static/admin.js +++ b/turnstone/console/static/admin.js @@ -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", diff --git a/turnstone/core/bm25.py b/turnstone/core/bm25.py index 74b2711a..7b8250b5 100644 --- a/turnstone/core/bm25.py +++ b/turnstone/core/bm25.py @@ -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) diff --git a/turnstone/core/memory_relevance.py b/turnstone/core/memory_relevance.py index d843f927..ef91ea93 100644 --- a/turnstone/core/memory_relevance.py +++ b/turnstone/core/memory_relevance.py @@ -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] diff --git a/turnstone/core/rerank.py b/turnstone/core/rerank.py index 990b158a..29a2889f 100644 --- a/turnstone/core/rerank.py +++ b/turnstone/core/rerank.py @@ -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: diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 180af7ce..57615a43 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -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] diff --git a/turnstone/core/settings_registry.py b/turnstone/core/settings_registry.py index d82aae6a..b6bead38 100644 --- a/turnstone/core/settings_registry.py +++ b/turnstone/core/settings_registry.py @@ -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", diff --git a/turnstone/core/tool_search.py b/turnstone/core/tool_search.py index 14c394d1..bc4b662a 100644 --- a/turnstone/core/tool_search.py +++ b/turnstone/core/tool_search.py @@ -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) diff --git a/turnstone/core/web_search.py b/turnstone/core/web_search.py index b28f2ff1..50cac666 100644 --- a/turnstone/core/web_search.py +++ b/turnstone/core/web_search.py @@ -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