mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
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:
@@ -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.
|
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):
|
Example (vLLM serving a Qwen3 reranker):
|
||||||
|
|
||||||
```toml
|
```toml
|
||||||
|
|||||||
@@ -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
@@ -1,6 +1,6 @@
|
|||||||
"""Tests for turnstone.core.bm25 — tokenizer and BM25 index."""
|
"""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:
|
class TestTokenize:
|
||||||
@@ -69,3 +69,179 @@ class TestBM25Index:
|
|||||||
results = index.search("file operations", k=3)
|
results = index.search("file operations", k=3)
|
||||||
# Doc 2 has more file/operations mentions, should rank higher
|
# Doc 2 has more file/operations mentions, should rank higher
|
||||||
assert results[0] == 2
|
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
|
||||||
|
|||||||
@@ -75,6 +75,73 @@ class TestScoreMemories:
|
|||||||
assert result[0]["name"] == "x"
|
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
|
# build_memory_context
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -395,6 +462,47 @@ class TestCompositionCandidateSelection:
|
|||||||
assert search_mock.call_args.args[1] == [("coordinator", "coord-1")]
|
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:
|
class TestMemorySearchToolExecution:
|
||||||
"""End-to-end test of ``memory(action='search')`` through _exec_memory.
|
"""End-to-end test of ``memory(action='search')`` through _exec_memory.
|
||||||
|
|
||||||
|
|||||||
@@ -301,3 +301,115 @@ class TestSessionRerankWiring:
|
|||||||
_registry=SimpleNamespace(get_config=lambda a: cfg),
|
_registry=SimpleNamespace(get_config=lambda a: cfg),
|
||||||
)
|
)
|
||||||
assert ChatSession._resolve_rerank_client(stub) is None # falls through, no rerank_url
|
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
|
||||||
|
|||||||
@@ -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.
|
# Truncation budget — required by _truncate_output on every exec.
|
||||||
session.tool_truncation = 100_000
|
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
|
# set_skill stub for load action. Records both the name and the
|
||||||
# arguments-string so tests can assert that #572's invocation-args
|
# arguments-string so tests can assert that #572's invocation-args
|
||||||
|
|||||||
@@ -218,6 +218,39 @@ class TestToolSearchManager:
|
|||||||
assert "mcp__github__create_issue" in text
|
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
|
# Helper function tests
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|||||||
@@ -5205,7 +5205,7 @@ const MODEL_ROLES = [
|
|||||||
{
|
{
|
||||||
label: "Reranker",
|
label: "Reranker",
|
||||||
description:
|
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",
|
aliasKey: "tools.reranker_alias",
|
||||||
fallbackKind: "disabled",
|
fallbackKind: "disabled",
|
||||||
mediaCapability: "supports_rerank",
|
mediaCapability: "supports_rerank",
|
||||||
|
|||||||
+75
-4
@@ -8,9 +8,19 @@ from __future__ import annotations
|
|||||||
import math
|
import math
|
||||||
import re
|
import re
|
||||||
from collections import Counter
|
from collections import Counter
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from turnstone.core.rerank import Reranker
|
||||||
|
|
||||||
_SPLIT_RE = re.compile(r"[_\-./\s]+")
|
_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]:
|
def _tokenize(text: str) -> list[str]:
|
||||||
"""Split text on whitespace, underscores, hyphens, dots."""
|
"""Split text on whitespace, underscores, hyphens, dots."""
|
||||||
@@ -20,9 +30,19 @@ def _tokenize(text: str) -> list[str]:
|
|||||||
class BM25Index:
|
class BM25Index:
|
||||||
"""Okapi BM25 ranking index over short text documents."""
|
"""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.k1 = k1
|
||||||
self.b = b
|
self.b = b
|
||||||
|
self._reranker = reranker
|
||||||
|
self._rerank_filters = rerank_filters
|
||||||
self._docs = documents
|
self._docs = documents
|
||||||
self._doc_tokens: list[list[str]] = [_tokenize(d) for d in documents]
|
self._doc_tokens: list[list[str]] = [_tokenize(d) for d in documents]
|
||||||
self._doc_lens = [len(t) for t in self._doc_tokens]
|
self._doc_lens = [len(t) for t in self._doc_tokens]
|
||||||
@@ -34,8 +54,8 @@ class BM25Index:
|
|||||||
for term in set(tokens):
|
for term in set(tokens):
|
||||||
self._df[term] += 1
|
self._df[term] += 1
|
||||||
|
|
||||||
def search(self, query: str, k: int = 5) -> list[int]:
|
def _bm25_rank(self, query: str) -> list[int]:
|
||||||
"""Return indices of top-k documents sorted by descending BM25 score."""
|
"""Return ALL matching document indices sorted by descending BM25 score."""
|
||||||
q_tokens = _tokenize(query)
|
q_tokens = _tokenize(query)
|
||||||
if not q_tokens:
|
if not q_tokens:
|
||||||
return []
|
return []
|
||||||
@@ -45,7 +65,58 @@ class BM25Index:
|
|||||||
if score > 0:
|
if score > 0:
|
||||||
scores.append((score, idx))
|
scores.append((score, idx))
|
||||||
scores.sort(key=lambda x: (-x[0], x[1]))
|
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:
|
def _score(self, q_tokens: list[str], doc_tokens: list[str], dl: int) -> float:
|
||||||
tf_map: Counter[str] = Counter(doc_tokens)
|
tf_map: Counter[str] = Counter(doc_tokens)
|
||||||
|
|||||||
@@ -4,10 +4,13 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from html import escape as _html_escape
|
from html import escape as _html_escape
|
||||||
from typing import Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from turnstone.core.bm25 import BM25Index
|
from turnstone.core.bm25 import BM25Index
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from turnstone.core.rerank import Reranker
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class MemoryConfig:
|
class MemoryConfig:
|
||||||
@@ -24,6 +27,8 @@ def score_memories(
|
|||||||
memories: list[dict[str, str]],
|
memories: list[dict[str, str]],
|
||||||
query: str,
|
query: str,
|
||||||
k: int = 5,
|
k: int = 5,
|
||||||
|
reranker: Reranker | None = None,
|
||||||
|
rerank_filters: bool = False,
|
||||||
) -> list[dict[str, str]]:
|
) -> list[dict[str, str]]:
|
||||||
"""Return the top-k memories most relevant to *query*.
|
"""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]}"
|
f"{m.get('name', '')} {m.get('description', '')} {m.get('content', '')[:200]}"
|
||||||
for m in memories
|
for m in memories
|
||||||
]
|
]
|
||||||
index = BM25Index(documents)
|
index = BM25Index(documents, reranker=reranker, rerank_filters=rerank_filters)
|
||||||
top_indices = index.search(query, k)
|
top_indices = index.search(query, k)
|
||||||
return [memories[i] for i in top_indices]
|
return [memories[i] for i in top_indices]
|
||||||
|
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ default and no fall back to a local model.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from collections.abc import Callable
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any, Protocol
|
from typing import Any, Protocol
|
||||||
|
|
||||||
@@ -30,6 +31,22 @@ from turnstone.core.log import get_logger
|
|||||||
|
|
||||||
log = get_logger(__name__)
|
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)
|
@dataclass(frozen=True)
|
||||||
class RerankHit:
|
class RerankHit:
|
||||||
|
|||||||
@@ -142,9 +142,9 @@ if TYPE_CHECKING:
|
|||||||
ModelCapabilities,
|
ModelCapabilities,
|
||||||
StreamChunk,
|
StreamChunk,
|
||||||
)
|
)
|
||||||
from turnstone.core.rerank import RerankClient
|
from turnstone.core.rerank import RerankClient, Reranker
|
||||||
from turnstone.core.tool_advisory import ToolAdvisory
|
from turnstone.core.tool_advisory import ToolAdvisory
|
||||||
from turnstone.core.web_search import Reranker, WebSearchClient
|
from turnstone.core.web_search import WebSearchClient
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# Cancellation support
|
# Cancellation support
|
||||||
@@ -604,6 +604,10 @@ _WATCH_QUEUE_SOFT_CAP = 50
|
|||||||
# the codepoint count.
|
# the codepoint count.
|
||||||
REMINDER_TEXT_STORAGE_CAP = 8192
|
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]]:
|
def _without_tool(tools: list[dict[str, Any]], name: str) -> list[dict[str, Any]]:
|
||||||
"""Return *tools* with the named tool removed."""
|
"""Return *tools* with the named tool removed."""
|
||||||
@@ -1178,6 +1182,7 @@ class ChatSession:
|
|||||||
self._tools,
|
self._tools,
|
||||||
always_on_names=builtin_in_session,
|
always_on_names=builtin_in_session,
|
||||||
max_results=tool_search_max_results,
|
max_results=tool_search_max_results,
|
||||||
|
reranker=self._bm25_reranker(),
|
||||||
)
|
)
|
||||||
# Skill: explicit name overrides is_default skills. ``skill_arguments``
|
# Skill: explicit name overrides is_default skills. ``skill_arguments``
|
||||||
# carries the spec's $ARGUMENTS payload — set at create/load time,
|
# carries the spec's $ARGUMENTS payload — set at create/load time,
|
||||||
@@ -1339,7 +1344,7 @@ class ChatSession:
|
|||||||
url=cfg.base_url,
|
url=cfg.base_url,
|
||||||
model=cfg.model or "",
|
model=cfg.model or "",
|
||||||
api_key=cfg.api_key,
|
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:
|
def _setting(key: str, env_value: str) -> str:
|
||||||
@@ -1355,11 +1360,11 @@ class ChatSession:
|
|||||||
url=_setting("tools.rerank_url", get_rerank_url()),
|
url=_setting("tools.rerank_url", get_rerank_url()),
|
||||||
model=_setting("tools.rerank_model", get_rerank_model()),
|
model=_setting("tools.rerank_model", get_rerank_model()),
|
||||||
api_key=_setting("tools.rerank_api_key", get_rerank_api_key()),
|
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:
|
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
|
The per-tool toggles default on; the operative gate is whether an
|
||||||
endpoint is configured (``_resolve_rerank_client`` returns None when
|
endpoint is configured (``_resolve_rerank_client`` returns None when
|
||||||
@@ -1387,6 +1392,51 @@ class ChatSession:
|
|||||||
|
|
||||||
return _rank
|
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(
|
def _resolve_capabilities(
|
||||||
self,
|
self,
|
||||||
provider: LLMProvider,
|
provider: LLMProvider,
|
||||||
@@ -2064,6 +2114,7 @@ class ChatSession:
|
|||||||
self._tools,
|
self._tools,
|
||||||
always_on_names=set(BUILTIN_TOOL_NAMES),
|
always_on_names=set(BUILTIN_TOOL_NAMES),
|
||||||
max_results=self._tool_search_max_results,
|
max_results=self._tool_search_max_results,
|
||||||
|
reranker=self._bm25_reranker(),
|
||||||
)
|
)
|
||||||
# Restore previously expanded tools that still exist
|
# Restore previously expanded tools that still exist
|
||||||
if old_expanded:
|
if old_expanded:
|
||||||
@@ -2767,7 +2818,14 @@ class ChatSession:
|
|||||||
context = extract_recent_context(self.messages)
|
context = extract_recent_context(self.messages)
|
||||||
visible_mems, candidate_source = self._select_memory_candidates(context)
|
visible_mems, candidate_source = self._select_memory_candidates(context)
|
||||||
if visible_mems:
|
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(
|
log.info(
|
||||||
"memory.composition",
|
"memory.composition",
|
||||||
source=candidate_source,
|
source=candidate_source,
|
||||||
@@ -8735,7 +8793,7 @@ class ChatSession:
|
|||||||
)
|
)
|
||||||
for r in rows
|
for r in rows
|
||||||
]
|
]
|
||||||
index = BM25Index(corpus)
|
index = BM25Index(corpus, reranker=self._bm25_reranker())
|
||||||
top = index.search(query, k=min(len(rows), 50))
|
top = index.search(query, k=min(len(rows), 50))
|
||||||
rows = [rows[i] for i in top]
|
rows = [rows[i] for i in top]
|
||||||
skills = [self._skills_project_row(r) for r in rows]
|
skills = [self._skills_project_row(r) for r in rows]
|
||||||
|
|||||||
@@ -307,6 +307,30 @@ def _build_registry() -> dict[str, SettingDef]:
|
|||||||
"compatible /rerank endpoint, then pick it under Models -> Roles -> Reranker. Takes "
|
"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).",
|
"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 ---------------------------------------------------------
|
# -- server ---------------------------------------------------------
|
||||||
SettingDef(
|
SettingDef(
|
||||||
"server.workstream_idle_timeout",
|
"server.workstream_idle_timeout",
|
||||||
|
|||||||
@@ -10,10 +10,13 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import re
|
import re
|
||||||
from collections import Counter
|
from collections import Counter
|
||||||
from typing import Any
|
from typing import TYPE_CHECKING, Any
|
||||||
|
|
||||||
from turnstone.core.bm25 import BM25Index, _tokenize # noqa: F401
|
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
|
# Tool search manager — partitions tools, tracks visibility
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
@@ -65,6 +68,7 @@ class ToolSearchManager:
|
|||||||
always_on_names: set[str],
|
always_on_names: set[str],
|
||||||
*,
|
*,
|
||||||
max_results: int = 5,
|
max_results: int = 5,
|
||||||
|
reranker: Reranker | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._always_on: list[dict[str, Any]] = []
|
self._always_on: list[dict[str, Any]] = []
|
||||||
self._deferred: list[dict[str, Any]] = []
|
self._deferred: list[dict[str, Any]] = []
|
||||||
@@ -82,7 +86,7 @@ class ToolSearchManager:
|
|||||||
|
|
||||||
# BM25 index over deferred tools
|
# BM25 index over deferred tools
|
||||||
texts = [_tool_text(t) for t in self._deferred]
|
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
|
# Pre-compute server summary for the search tool description
|
||||||
self._server_hint = _mcp_server_summary(self._deferred)
|
self._server_hint = _mcp_server_summary(self._deferred)
|
||||||
|
|||||||
@@ -16,7 +16,6 @@ clients.
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from collections.abc import Callable
|
|
||||||
from typing import TYPE_CHECKING, Any, Protocol
|
from typing import TYPE_CHECKING, Any, Protocol
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
@@ -25,16 +24,13 @@ from turnstone.core.log import get_logger
|
|||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from turnstone.core.mcp_client import MCPClientManager
|
from turnstone.core.mcp_client import MCPClientManager
|
||||||
|
from turnstone.core.rerank import Reranker
|
||||||
|
|
||||||
log = get_logger(__name__)
|
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
|
# 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
|
_RERANK_POOL = 50
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user