mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
feat(rerank): endpoint-backed reranking for web_search
Reranking is delegated to an external Cohere/Jina-compatible /rerank endpoint (self-hosted vLLM/TEI/llama.cpp, or hosted Cohere/Jina/Voyage); Turnstone runs no reranker model itself. Disabled until an endpoint is configured. - core/rerank.py: CohereJinaRerankClient (tolerant of results-wrapped and bare-list responses) + resolver. - web_search: rerank the SearxNG result pool by query relevance before top-k, with a native-order fallback on error; answers/infoboxes untouched. - Reranker as a model definition: add a model with the supports_rerank capability and pick it under Models -> Roles -> Reranker (tools.reranker_alias); takes precedence over the tools.rerank_url settings. Settings: tools.rerank_url/model/api_key, tools.rerank_web_search, tools.reranker_alias. Docs: docs/tools.md, turnstone.example.toml. (web_fetch reranking was evaluated and dropped: for single-document chunk selection it did not reliably beat head-truncation. Reranking is reserved for multi-item ranking.)
This commit is contained in:
Patrick Buckley
@@ -310,6 +310,27 @@ Search the web using a text query.
|
||||
|
||||
---
|
||||
|
||||
### Reranking (optional)
|
||||
|
||||
`web_search` can use an external **reranker** to re-order the backend's result pool by relevance to the query before returning the top hits. Turnstone runs no reranker model itself; it POSTs to a Cohere/Jina-compatible `/rerank` endpoint (self-hosted [vLLM](https://docs.vllm.ai) / [TEI](https://github.com/huggingface/text-embeddings-inference) / llama.cpp, or hosted Cohere/Jina/Voyage).
|
||||
|
||||
**Disabled by default.** Two ways to point at an endpoint:
|
||||
|
||||
- **A reranker model (recommended, admin UI):** in the console **Models** tab, add a model definition whose `base_url` is a Cohere/Jina-compatible `/rerank` endpoint and whose capabilities include `{"supports_rerank": true}`, then select it under **Models → Roles → Reranker**. Managed like every other model (write-only key, enable/disable).
|
||||
- **Config settings:** set `rerank_url` (full endpoint URL including path) in `config.toml` `[tools]`, the admin Settings tab, or `$TURNSTONE_RERANK_URL`; optionally `rerank_model` and `rerank_api_key` (write-only). The reranker model role, if set, takes precedence over these.
|
||||
|
||||
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.
|
||||
|
||||
Example (vLLM serving a Qwen3 reranker):
|
||||
|
||||
```toml
|
||||
[tools]
|
||||
rerank_url = "http://vllm:8000/rerank"
|
||||
rerank_model = "qwen3-reranker"
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Agent
|
||||
|
||||
The tool name uses the `_agent` suffix — bare `task` collides with
|
||||
|
||||
@@ -0,0 +1,303 @@
|
||||
"""Tests for turnstone.core.rerank — endpoint-backed reranking client."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from turnstone.core.rerank import (
|
||||
CohereJinaRerankClient,
|
||||
RerankHit,
|
||||
_parse_hits,
|
||||
resolve_rerank_client,
|
||||
)
|
||||
from turnstone.core.session import ChatSession
|
||||
|
||||
|
||||
def _mock_httpx_post(handler):
|
||||
"""Patch target for ``rerank.httpx.post`` that routes the call through a real
|
||||
``httpx.MockTransport``. The request flows through genuine httpx JSON/header
|
||||
encoding and response parsing — a true boundary, not a bare MagicMock.
|
||||
"""
|
||||
client = httpx.Client(transport=httpx.MockTransport(handler))
|
||||
|
||||
def _post(url, **kwargs):
|
||||
return client.post(url, **kwargs)
|
||||
|
||||
return _post
|
||||
|
||||
|
||||
# A Cohere/Jina/vLLM-shaped response: results wrapper + relevance_score, returned
|
||||
# out of input order so tests prove the client sorts best-first.
|
||||
RESULTS_WRAPPED = {
|
||||
"results": [
|
||||
{"index": 2, "relevance_score": 0.10},
|
||||
{"index": 0, "relevance_score": 0.95},
|
||||
{"index": 1, "relevance_score": 0.42},
|
||||
]
|
||||
}
|
||||
|
||||
# A TEI-shaped response: bare list + "score" key.
|
||||
BARE_LIST = [
|
||||
{"index": 0, "score": 0.30},
|
||||
{"index": 1, "score": 0.80},
|
||||
]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _parse_hits — response-shape tolerance (the boundary that varies by provider)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestParseHits:
|
||||
def test_results_wrapped_relevance_score_sorted(self):
|
||||
hits = _parse_hits(RESULTS_WRAPPED, n_docs=3)
|
||||
assert [(h.index, h.score) for h in hits] == [(0, 0.95), (1, 0.42), (2, 0.10)]
|
||||
|
||||
def test_bare_list_score_key(self):
|
||||
hits = _parse_hits(BARE_LIST, n_docs=2)
|
||||
assert [(h.index, h.score) for h in hits] == [(1, 0.80), (0, 0.30)]
|
||||
|
||||
def test_relevance_score_preferred_over_score(self):
|
||||
# When both keys are present, relevance_score wins.
|
||||
hits = _parse_hits({"results": [{"index": 0, "relevance_score": 0.9, "score": 0.1}]}, 1)
|
||||
assert hits == [RerankHit(index=0, score=0.9)]
|
||||
|
||||
def test_drops_out_of_range_index(self):
|
||||
hits = _parse_hits({"results": [{"index": 5, "relevance_score": 0.9}]}, n_docs=2)
|
||||
assert hits == []
|
||||
|
||||
def test_drops_non_dict_and_missing_fields(self):
|
||||
data = {
|
||||
"results": [
|
||||
"not-a-dict",
|
||||
{"index": 0}, # missing score
|
||||
{"relevance_score": 0.5}, # missing index
|
||||
{"index": 1, "relevance_score": 0.7}, # valid
|
||||
]
|
||||
}
|
||||
assert _parse_hits(data, n_docs=2) == [RerankHit(index=1, score=0.7)]
|
||||
|
||||
def test_rejects_bool_index_and_score(self):
|
||||
# bool is a subclass of int/float — must not be accepted as a hit.
|
||||
assert _parse_hits({"results": [{"index": True, "relevance_score": 0.9}]}, 2) == []
|
||||
assert _parse_hits({"results": [{"index": 0, "relevance_score": True}]}, 2) == []
|
||||
|
||||
def test_empty_and_garbage(self):
|
||||
assert _parse_hits({"results": []}, 3) == []
|
||||
assert _parse_hits({}, 3) == []
|
||||
assert _parse_hits({"results": "nope"}, 3) == []
|
||||
assert _parse_hits(42, 3) == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CohereJinaRerankClient — request construction + transport boundary
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCohereJinaRerankClient:
|
||||
def test_request_body_boundary(self, monkeypatch):
|
||||
"""Drive a real httpx request through MockTransport and assert the URL,
|
||||
body, and auth header the client builds."""
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["url"] = str(request.url)
|
||||
captured["body"] = json.loads(request.content)
|
||||
captured["auth"] = request.headers.get("authorization")
|
||||
return httpx.Response(200, json=RESULTS_WRAPPED)
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
client = CohereJinaRerankClient(
|
||||
"http://vllm:8000/rerank", model="bge", api_key="secret", timeout=10
|
||||
)
|
||||
hits = client.rerank("q", ["a", "b", "c"], top_n=2)
|
||||
|
||||
assert captured["url"] == "http://vllm:8000/rerank"
|
||||
assert captured["body"] == {
|
||||
"query": "q",
|
||||
"documents": ["a", "b", "c"],
|
||||
"model": "bge",
|
||||
"top_n": 2,
|
||||
}
|
||||
assert captured["auth"] == "Bearer secret"
|
||||
# Response is parsed + sorted best-first.
|
||||
assert [h.index for h in hits] == [0, 1, 2]
|
||||
|
||||
def test_model_and_top_n_omitted_when_unset(self, monkeypatch):
|
||||
captured: dict = {}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
captured["body"] = json.loads(request.content)
|
||||
captured["auth"] = request.headers.get("authorization")
|
||||
return httpx.Response(200, json={"results": []})
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
CohereJinaRerankClient("http://h/rerank").rerank("q", ["a"])
|
||||
|
||||
assert captured["body"] == {"query": "q", "documents": ["a"]}
|
||||
assert captured["auth"] is None # no Authorization header without a key
|
||||
|
||||
def test_empty_documents_makes_no_request(self, monkeypatch):
|
||||
called = {"n": 0}
|
||||
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
called["n"] += 1
|
||||
return httpx.Response(200, json={"results": []})
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
assert CohereJinaRerankClient("http://h/rerank").rerank("q", []) == []
|
||||
assert called["n"] == 0 # short-circuits before any HTTP call
|
||||
|
||||
def test_bare_list_response_parsed(self, monkeypatch):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json=BARE_LIST)
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
hits = CohereJinaRerankClient("http://tei/rerank").rerank("q", ["a", "b"])
|
||||
assert [h.index for h in hits] == [1, 0]
|
||||
|
||||
def test_http_error_propagates(self, monkeypatch):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(401, text="unauthorized")
|
||||
|
||||
monkeypatch.setattr("turnstone.core.rerank.httpx.post", _mock_httpx_post(handler))
|
||||
with pytest.raises(httpx.HTTPStatusError):
|
||||
CohereJinaRerankClient("http://h/rerank").rerank("q", ["a"])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# resolve_rerank_client
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestResolveRerankClient:
|
||||
def test_client_when_url_set(self):
|
||||
client = resolve_rerank_client("http://h/rerank", model="m", api_key="k")
|
||||
assert isinstance(client, CohereJinaRerankClient)
|
||||
assert client._url == "http://h/rerank"
|
||||
assert client._model == "m"
|
||||
assert client._api_key == "k"
|
||||
|
||||
def test_none_when_no_url(self):
|
||||
assert resolve_rerank_client("") is None
|
||||
assert resolve_rerank_client(" ") is None
|
||||
assert resolve_rerank_client(None) is None # type: ignore[arg-type]
|
||||
|
||||
def test_strips_whitespace(self):
|
||||
client = resolve_rerank_client(" http://h/rerank ", model=" m ", api_key=" k ")
|
||||
assert isinstance(client, CohereJinaRerankClient)
|
||||
assert client._url == "http://h/rerank"
|
||||
assert client._model == "m"
|
||||
assert client._api_key == "k"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ChatSession wiring — disabled by default, endpoint-gated, per-tool toggles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeConfigStore:
|
||||
"""Minimal ConfigStore stand-in (explicit value vs unset via stored_keys)."""
|
||||
|
||||
def __init__(self, values: dict, stored: set[str] | None = None) -> None:
|
||||
self._values = dict(values)
|
||||
self._stored = frozenset(stored if stored is not None else values.keys())
|
||||
|
||||
def stored_keys(self) -> frozenset[str]:
|
||||
return self._stored
|
||||
|
||||
def get(self, key: str):
|
||||
return self._values.get(key)
|
||||
|
||||
|
||||
def _patch_getters_empty(monkeypatch):
|
||||
for name in ("get_rerank_url", "get_rerank_model", "get_rerank_api_key"):
|
||||
monkeypatch.setattr(f"turnstone.core.config.{name}", lambda: "")
|
||||
|
||||
|
||||
class TestSessionRerankWiring:
|
||||
def test_disabled_by_default_when_no_endpoint(self, monkeypatch):
|
||||
_patch_getters_empty(monkeypatch)
|
||||
stub = SimpleNamespace(_config_store=None, tool_timeout=30)
|
||||
assert ChatSession._resolve_rerank_client(stub) is None
|
||||
|
||||
def test_resolves_client_from_config_store(self, monkeypatch):
|
||||
_patch_getters_empty(monkeypatch)
|
||||
cs = _FakeConfigStore(
|
||||
{
|
||||
"tools.rerank_url": "http://h/rerank",
|
||||
"tools.rerank_model": "m",
|
||||
"tools.rerank_api_key": "k",
|
||||
}
|
||||
)
|
||||
stub = SimpleNamespace(_config_store=cs, tool_timeout=15)
|
||||
client = ChatSession._resolve_rerank_client(stub)
|
||||
assert isinstance(client, CohereJinaRerankClient)
|
||||
assert client._url == "http://h/rerank"
|
||||
assert client._model == "m"
|
||||
assert client._api_key == "k"
|
||||
|
||||
def test_explicit_empty_url_stays_disabled(self, monkeypatch):
|
||||
# An admin who clears the URL (explicit "") must NOT fall back to a
|
||||
# config.toml/env value — explicit-empty means "off".
|
||||
monkeypatch.setattr("turnstone.core.config.get_rerank_url", lambda: "http://env/rerank")
|
||||
monkeypatch.setattr("turnstone.core.config.get_rerank_model", lambda: "")
|
||||
monkeypatch.setattr("turnstone.core.config.get_rerank_api_key", lambda: "")
|
||||
cs = _FakeConfigStore({"tools.rerank_url": ""}, stored={"tools.rerank_url"})
|
||||
stub = SimpleNamespace(_config_store=cs, tool_timeout=30)
|
||||
assert ChatSession._resolve_rerank_client(stub) is None
|
||||
|
||||
def test_enabled_for_defaults_true_without_store(self):
|
||||
stub = SimpleNamespace(_config_store=None)
|
||||
assert ChatSession._rerank_enabled_for(stub, "web_search") is True
|
||||
|
||||
def test_enabled_for_respects_per_tool_toggle(self):
|
||||
cs = _FakeConfigStore({"tools.rerank_web_search": False})
|
||||
stub = SimpleNamespace(_config_store=cs)
|
||||
assert ChatSession._rerank_enabled_for(stub, "web_search") is False
|
||||
|
||||
def test_prefers_reranker_model_definition(self, monkeypatch):
|
||||
# A model definition with supports_rerank, selected via the Reranker
|
||||
# role, wins over the raw tools.rerank_url settings.
|
||||
_patch_getters_empty(monkeypatch)
|
||||
from turnstone.core.model_registry import ModelConfig
|
||||
|
||||
cfg = ModelConfig(
|
||||
alias="rr",
|
||||
base_url="http://rr:8000/rerank",
|
||||
api_key="k",
|
||||
model="bge",
|
||||
capabilities={"supports_rerank": True},
|
||||
)
|
||||
cs = _FakeConfigStore({"tools.reranker_alias": "rr"})
|
||||
stub = SimpleNamespace(
|
||||
_config_store=cs,
|
||||
tool_timeout=30,
|
||||
_registry=SimpleNamespace(get_config=lambda a: cfg),
|
||||
)
|
||||
client = ChatSession._resolve_rerank_client(stub)
|
||||
assert isinstance(client, CohereJinaRerankClient)
|
||||
assert client._url == "http://rr:8000/rerank"
|
||||
assert client._model == "bge"
|
||||
assert client._api_key == "k"
|
||||
|
||||
def test_ignores_alias_without_rerank_capability(self, monkeypatch):
|
||||
# A non-reranker model (no supports_rerank) must NOT be used as a reranker,
|
||||
# even if the alias is set — guards against a stale/wrong alias misrouting.
|
||||
_patch_getters_empty(monkeypatch)
|
||||
from turnstone.core.model_registry import ModelConfig
|
||||
|
||||
cfg = ModelConfig(
|
||||
alias="chat", base_url="http://chat/v1", api_key="k", model="gpt", capabilities={}
|
||||
)
|
||||
cs = _FakeConfigStore({"tools.reranker_alias": "chat"})
|
||||
stub = SimpleNamespace(
|
||||
_config_store=cs,
|
||||
tool_timeout=30,
|
||||
_registry=SimpleNamespace(get_config=lambda a: cfg),
|
||||
)
|
||||
assert ChatSession._resolve_rerank_client(stub) is None # falls through, no rerank_url
|
||||
@@ -317,3 +317,94 @@ class TestResolveClient:
|
||||
"mcp:static-search:search", searxng_url=None, mcp_client=mcp
|
||||
)
|
||||
assert isinstance(client, MCPSearchClient)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reranking
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _results(*titles):
|
||||
return {
|
||||
"results": [
|
||||
{"title": t, "url": f"http://{t.lower()}", "content": f"{t} snippet"} for t in titles
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
class TestWebSearchReranking:
|
||||
def test_reorders_results_by_reranker_output(self):
|
||||
# Reranker promotes index 2, then 0, then 1.
|
||||
out = _format_searxng(
|
||||
_results("A", "B", "C"), "q", max_results=3, reranker=lambda q, d: [2, 0, 1]
|
||||
)
|
||||
assert out.index("[C]") < out.index("[A]") < out.index("[B]")
|
||||
|
||||
def test_reranker_receives_title_and_snippet(self):
|
||||
seen: dict = {}
|
||||
|
||||
def rr(query, docs):
|
||||
seen["query"] = query
|
||||
seen["docs"] = docs
|
||||
return list(range(len(docs)))
|
||||
|
||||
_format_searxng(_results("Py", "X"), "find me", reranker=rr)
|
||||
assert seen["query"] == "find me"
|
||||
assert seen["docs"][0] == "Py\nPy snippet"
|
||||
|
||||
def test_error_falls_back_to_native_order(self):
|
||||
def boom(query, docs):
|
||||
raise RuntimeError("rerank endpoint down")
|
||||
|
||||
out = _format_searxng(_results("A", "B"), "q", reranker=boom)
|
||||
assert out.index("[A]") < out.index("[B]") # native order preserved
|
||||
|
||||
def test_skipped_for_single_result(self):
|
||||
called = {"n": 0}
|
||||
|
||||
def rr(query, docs):
|
||||
called["n"] += 1
|
||||
return [0]
|
||||
|
||||
_format_searxng(_results("Only"), "q", reranker=rr)
|
||||
assert called["n"] == 0 # <=1 result: nothing to reorder
|
||||
|
||||
def test_answers_and_infoboxes_untouched(self):
|
||||
# Reranking only reorders the results list, never answers/infoboxes.
|
||||
out = _format_searxng(SEARXNG_JSON, "python", reranker=lambda q, d: [1, 0])
|
||||
assert "Answer: Python is a programming language." in out
|
||||
assert "A high-level language." in out
|
||||
|
||||
def test_partial_order_keeps_all_results(self):
|
||||
# A top_n-style reranker returns only a subset; the rest must survive.
|
||||
out = _format_searxng(
|
||||
_results("T0", "T1", "T2"), "q", max_results=3, reranker=lambda q, d: [1]
|
||||
)
|
||||
assert "[T0]" in out and "[T1]" in out and "[T2]" in out
|
||||
assert out.index("[T1]") < out.index("[T0]") # T1 promoted
|
||||
|
||||
def test_searxng_client_threads_reranker_kwarg(self):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, json=_results("A", "B"))
|
||||
|
||||
with patch("turnstone.core.web_search.httpx.get", _mock_httpx_get(handler)):
|
||||
out = SearXNGClient("http://searxng:8080").search("q", reranker=lambda q, d: [1, 0])
|
||||
|
||||
assert out.index("[B]") < out.index("[A]") # reranker applied via search()
|
||||
|
||||
def test_pool_cap_preserves_tail_beyond_50(self):
|
||||
# >_RERANK_POOL (50) results: only the first 50 are reranked; the tail
|
||||
# must survive, appended in native order after the reranked pool.
|
||||
data = {
|
||||
"results": [
|
||||
{"title": f"R{i}", "url": f"http://r/{i}", "content": f"c{i}"} for i in range(60)
|
||||
]
|
||||
}
|
||||
# Reranker reverses the 50-item pool it is handed.
|
||||
out = _format_searxng(
|
||||
data, "q", max_results=60, reranker=lambda q, d: list(range(len(d)))[::-1]
|
||||
)
|
||||
assert all(f"[R{i}]" in out for i in range(60)) # nothing dropped
|
||||
assert out.index("[R49]") < out.index("[R0]") # pool reversed
|
||||
assert out.index("[R0]") < out.index("[R50]") # reranked pool before the tail
|
||||
assert out.index("[R50]") < out.index("[R59]") # tail kept in native order
|
||||
|
||||
@@ -109,6 +109,22 @@
|
||||
# searxng_engines = "" # comma-separated engines (e.g. "duckduckgo,wikipedia");
|
||||
# empty = the instance's default mix
|
||||
# env: TURNSTONE_SEARXNG_ENGINES
|
||||
#
|
||||
# Reranking (optional, disabled by default). Turnstone runs no reranker itself —
|
||||
# it POSTs to an external Cohere/Jina-compatible /rerank endpoint (self-hosted
|
||||
# vLLM/TEI/llama.cpp, or hosted Cohere/Jina/Voyage) to reorder web_search results
|
||||
# by query relevance. Set rerank_url to enable; leave empty to disable. (Alternatively,
|
||||
# add a reranker model in the admin Models tab and pick it under Models -> Roles
|
||||
# -> Reranker; that selection takes precedence over rerank_url.)
|
||||
# rerank_url = "" # full endpoint URL incl. path, e.g.
|
||||
# "http://vllm:8000/rerank" or
|
||||
# "https://api.cohere.com/v2/rerank";
|
||||
# env: TURNSTONE_RERANK_URL
|
||||
# rerank_model = "" # model name sent in the request; empty = the
|
||||
# endpoint's default; env: TURNSTONE_RERANK_MODEL
|
||||
# rerank_api_key = "" # bearer token for hosted providers; usually empty
|
||||
# for self-hosted. config.toml only (never env)
|
||||
# rerank_web_search = true # rerank web_search results (when an endpoint is set)
|
||||
|
||||
# --- Judge (turnstone, node) ---
|
||||
|
||||
|
||||
@@ -5202,6 +5202,15 @@ const MODEL_ROLES = [
|
||||
mediaCapability: "supports_speech_synthesis",
|
||||
mediaRole: "tts",
|
||||
},
|
||||
{
|
||||
label: "Reranker",
|
||||
description:
|
||||
"Reranks web_search results and selects the most relevant chunks of large web_fetch pages. 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).",
|
||||
aliasKey: "tools.reranker_alias",
|
||||
fallbackKind: "disabled",
|
||||
mediaCapability: "supports_rerank",
|
||||
mediaRole: "rerank",
|
||||
},
|
||||
];
|
||||
|
||||
// Whether a model definition is eligible for an audio role. Mirrors
|
||||
|
||||
@@ -229,6 +229,57 @@ def get_searxng_engines() -> str:
|
||||
return _searxng_engines
|
||||
|
||||
|
||||
# -- Rerank endpoint (cached) -------------------------------------------------
|
||||
|
||||
_rerank_url: str | None = None
|
||||
_rerank_url_loaded: bool = False
|
||||
_rerank_model: str | None = None
|
||||
_rerank_model_loaded: bool = False
|
||||
_rerank_api_key: str | None = None
|
||||
_rerank_api_key_loaded: bool = False
|
||||
|
||||
|
||||
def get_rerank_url() -> str:
|
||||
"""Load the rerank endpoint URL (cached after first call).
|
||||
|
||||
Precedence: config.toml [tools] rerank_url -> $TURNSTONE_RERANK_URL -> ""
|
||||
("" disables reranking — there is no bundled rerank endpoint).
|
||||
"""
|
||||
global _rerank_url, _rerank_url_loaded
|
||||
if not _rerank_url_loaded:
|
||||
_rerank_url_loaded = True
|
||||
cfg = load_config("tools").get("rerank_url", "").strip()
|
||||
_rerank_url = cfg or os.environ.get("TURNSTONE_RERANK_URL", "").strip()
|
||||
return _rerank_url or ""
|
||||
|
||||
|
||||
def get_rerank_model() -> str:
|
||||
"""Load the rerank model name (cached after first call).
|
||||
|
||||
Precedence: config.toml [tools] rerank_model -> $TURNSTONE_RERANK_MODEL -> ""
|
||||
(use the endpoint's default model).
|
||||
"""
|
||||
global _rerank_model, _rerank_model_loaded
|
||||
if not _rerank_model_loaded:
|
||||
_rerank_model_loaded = True
|
||||
cfg = load_config("tools").get("rerank_model", "").strip()
|
||||
_rerank_model = cfg or os.environ.get("TURNSTONE_RERANK_MODEL", "").strip()
|
||||
return _rerank_model or ""
|
||||
|
||||
|
||||
def get_rerank_api_key() -> str:
|
||||
"""Load the rerank endpoint bearer token (cached after first call).
|
||||
|
||||
Read from config.toml [tools] rerank_api_key only — secrets are never read
|
||||
from the environment.
|
||||
"""
|
||||
global _rerank_api_key, _rerank_api_key_loaded
|
||||
if not _rerank_api_key_loaded:
|
||||
_rerank_api_key_loaded = True
|
||||
_rerank_api_key = load_config("tools").get("rerank_api_key", "").strip()
|
||||
return _rerank_api_key or ""
|
||||
|
||||
|
||||
def nonneg_float(val: str) -> float:
|
||||
"""Argparse type for non-negative floats (``>= 0``)."""
|
||||
f = float(val)
|
||||
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Endpoint-backed reranking (Cohere/Jina-compatible wire format).
|
||||
|
||||
Turnstone performs no in-process model inference — reranking is delegated to an
|
||||
external rerank endpoint, exactly like every other model the platform talks to.
|
||||
The endpoint must speak the de-facto-standard Cohere/Jina ``/rerank`` contract,
|
||||
which is also implemented by self-hosted servers (vLLM, Text Embeddings
|
||||
Inference, llama.cpp) and the hosted Cohere / Jina / Voyage APIs.
|
||||
|
||||
Request (POST to the configured URL)::
|
||||
|
||||
{"model": "<name>", "query": "<q>", "documents": ["...", ...], "top_n": N}
|
||||
|
||||
Response — two shapes are accepted::
|
||||
|
||||
{"results": [{"index": 0, "relevance_score": 0.91}, ...]} # Cohere/Jina/vLLM
|
||||
[{"index": 0, "score": 0.91}, ...] # bare list (TEI)
|
||||
|
||||
Reranking is disabled unless an endpoint URL is configured; there is no bundled
|
||||
default and no fall back to a local model.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Protocol
|
||||
|
||||
import httpx
|
||||
|
||||
from turnstone.core.log import get_logger
|
||||
|
||||
log = get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RerankHit:
|
||||
"""One reranked document: its position in the input list and its score."""
|
||||
|
||||
index: int # 0-based index into the documents passed to ``rerank``
|
||||
score: float # relevance score; higher is more relevant
|
||||
|
||||
|
||||
class RerankClient(Protocol):
|
||||
"""Minimal interface for a rerank backend."""
|
||||
|
||||
def rerank(
|
||||
self, query: str, documents: list[str], *, top_n: int | None = None
|
||||
) -> list[RerankHit]:
|
||||
"""Score ``documents`` against ``query``; return hits sorted best-first."""
|
||||
...
|
||||
|
||||
|
||||
class CohereJinaRerankClient:
|
||||
"""Rerank via a Cohere/Jina-compatible ``POST <url>`` endpoint.
|
||||
|
||||
``url`` is the *full* endpoint (including path), because the path differs by
|
||||
provider — ``/rerank`` (vLLM, TEI), ``/v1/rerank`` (Jina, llama.cpp),
|
||||
``/v2/rerank`` (Cohere). The request body and the ``results`` /
|
||||
``relevance_score`` response are shared across all of them.
|
||||
"""
|
||||
|
||||
def __init__(self, url: str, model: str = "", api_key: str = "", timeout: float = 30) -> None:
|
||||
self._url = url
|
||||
self._model = model
|
||||
self._api_key = api_key
|
||||
self._timeout = timeout
|
||||
|
||||
def rerank(
|
||||
self, query: str, documents: list[str], *, top_n: int | None = None
|
||||
) -> list[RerankHit]:
|
||||
if not documents:
|
||||
return []
|
||||
payload: dict[str, Any] = {"query": query, "documents": list(documents)}
|
||||
if self._model:
|
||||
payload["model"] = self._model
|
||||
if top_n is not None:
|
||||
payload["top_n"] = top_n
|
||||
headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else {}
|
||||
resp = httpx.post(self._url, json=payload, headers=headers, timeout=self._timeout)
|
||||
resp.raise_for_status()
|
||||
return _parse_hits(resp.json(), len(documents))
|
||||
|
||||
|
||||
def _parse_hits(data: Any, n_docs: int) -> list[RerankHit]:
|
||||
"""Parse a Cohere/Jina ``{"results": [...]}`` or bare-list rerank response.
|
||||
|
||||
Tolerates both ``relevance_score`` (Cohere/Jina/vLLM) and ``score`` (TEI),
|
||||
drops malformed or out-of-range entries, and returns hits sorted best-first.
|
||||
"""
|
||||
rows = data.get("results") if isinstance(data, dict) else data
|
||||
if not isinstance(rows, list):
|
||||
return []
|
||||
hits: list[RerankHit] = []
|
||||
for row in rows:
|
||||
if not isinstance(row, dict):
|
||||
continue
|
||||
idx = row.get("index")
|
||||
score = row.get("relevance_score")
|
||||
if score is None:
|
||||
score = row.get("score")
|
||||
# bool is a subclass of int/float — reject it explicitly.
|
||||
if (
|
||||
isinstance(idx, int)
|
||||
and not isinstance(idx, bool)
|
||||
and 0 <= idx < n_docs
|
||||
and isinstance(score, (int, float))
|
||||
and not isinstance(score, bool)
|
||||
):
|
||||
hits.append(RerankHit(index=idx, score=float(score)))
|
||||
hits.sort(key=lambda h: h.score, reverse=True)
|
||||
return hits
|
||||
|
||||
|
||||
def resolve_rerank_client(
|
||||
url: str, model: str = "", api_key: str = "", timeout: float = 30
|
||||
) -> RerankClient | None:
|
||||
"""Return a rerank client, or ``None`` when no endpoint URL is configured.
|
||||
|
||||
A missing URL is the "reranking disabled" state (the default) — there is no
|
||||
bundled endpoint and no local-inference fallback.
|
||||
"""
|
||||
url = (url or "").strip()
|
||||
if not url:
|
||||
return None
|
||||
return CohereJinaRerankClient(
|
||||
url=url,
|
||||
model=(model or "").strip(),
|
||||
api_key=(api_key or "").strip(),
|
||||
timeout=timeout,
|
||||
)
|
||||
@@ -142,8 +142,9 @@ if TYPE_CHECKING:
|
||||
ModelCapabilities,
|
||||
StreamChunk,
|
||||
)
|
||||
from turnstone.core.rerank import RerankClient
|
||||
from turnstone.core.tool_advisory import ToolAdvisory
|
||||
from turnstone.core.web_search import WebSearchClient
|
||||
from turnstone.core.web_search import Reranker, WebSearchClient
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cancellation support
|
||||
@@ -1303,6 +1304,89 @@ class ChatSession:
|
||||
timeout=self.tool_timeout,
|
||||
)
|
||||
|
||||
def _resolve_rerank_client(self) -> RerankClient | None:
|
||||
"""Return a rerank client, or None when reranking is unconfigured.
|
||||
|
||||
Precedence: a reranker **model definition** selected via the Reranker
|
||||
role (``tools.reranker_alias`` → a model with ``supports_rerank``) wins;
|
||||
otherwise the ``tools.rerank_url`` settings (storage → config.toml/env).
|
||||
There is no bundled rerank endpoint, so reranking stays disabled until
|
||||
one is configured.
|
||||
"""
|
||||
from turnstone.core.config import (
|
||||
get_rerank_api_key,
|
||||
get_rerank_model,
|
||||
get_rerank_url,
|
||||
)
|
||||
from turnstone.core.rerank import resolve_rerank_client
|
||||
|
||||
cs = getattr(self, "_config_store", None)
|
||||
stored = cs.stored_keys() if cs is not None else frozenset()
|
||||
|
||||
# A reranker model definition (capability ``supports_rerank``), picked
|
||||
# via the Reranker role, takes precedence — managed like every other
|
||||
# model. Its base_url is the full Cohere/Jina-compatible rerank endpoint.
|
||||
registry = getattr(self, "_registry", None)
|
||||
if cs is not None and registry is not None:
|
||||
alias = str(cs.get("tools.reranker_alias") or "").strip()
|
||||
if alias:
|
||||
try:
|
||||
cfg = registry.get_config(alias)
|
||||
except Exception:
|
||||
cfg = None
|
||||
if cfg is not None and cfg.base_url and cfg.capabilities.get("supports_rerank"):
|
||||
return resolve_rerank_client(
|
||||
url=cfg.base_url,
|
||||
model=cfg.model or "",
|
||||
api_key=cfg.api_key,
|
||||
timeout=self.tool_timeout,
|
||||
)
|
||||
|
||||
def _setting(key: str, env_value: str) -> str:
|
||||
if cs is not None and key in stored: # explicit admin value wins
|
||||
return str(cs.get(key) or "").strip()
|
||||
if env_value: # config.toml / env var
|
||||
return env_value
|
||||
if cs is not None: # registry default, surfaced by the store
|
||||
return str(cs.get(key) or "").strip()
|
||||
return ""
|
||||
|
||||
return resolve_rerank_client(
|
||||
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,
|
||||
)
|
||||
|
||||
def _rerank_enabled_for(self, tool: str) -> bool:
|
||||
"""Whether reranking is enabled for ``tool`` ('web_search' | 'web_fetch').
|
||||
|
||||
The per-tool toggles default on; the operative gate is whether an
|
||||
endpoint is configured (``_resolve_rerank_client`` returns None when
|
||||
not). A bare CLI without a ConfigStore inherits the on-by-default toggle.
|
||||
"""
|
||||
cs = getattr(self, "_config_store", None)
|
||||
if cs is not None:
|
||||
return bool(cs.get(f"tools.rerank_{tool}"))
|
||||
return True
|
||||
|
||||
def _web_search_reranker(self) -> Reranker | None:
|
||||
"""Build a web_search reranker callable, or None when disabled.
|
||||
|
||||
Returns a ``(query, docs) -> ranked indices`` adapter over the configured
|
||||
rerank endpoint; None when reranking is off or no endpoint is set.
|
||||
"""
|
||||
if not self._rerank_enabled_for("web_search"):
|
||||
return None
|
||||
rc = self._resolve_rerank_client()
|
||||
if rc is None:
|
||||
return None
|
||||
|
||||
def _rank(query: str, docs: list[str]) -> list[int]:
|
||||
return [hit.index for hit in rc.rerank(query, docs)]
|
||||
|
||||
return _rank
|
||||
|
||||
def _resolve_capabilities(
|
||||
self,
|
||||
provider: LLMProvider,
|
||||
@@ -12130,7 +12214,12 @@ class ChatSession:
|
||||
return call_id, msg
|
||||
|
||||
try:
|
||||
output = client.search(query, max_results=max_results, category=category)
|
||||
output = client.search(
|
||||
query,
|
||||
max_results=max_results,
|
||||
category=category,
|
||||
reranker=self._web_search_reranker(),
|
||||
)
|
||||
except Exception as e:
|
||||
msg = f"Error: web search failed: {e}"
|
||||
self._report_tool_result(call_id, "web_search", msg, is_error=True)
|
||||
|
||||
@@ -250,6 +250,62 @@ def _build_registry() -> dict[str, SettingDef]:
|
||||
"'searxng' forces the SearxNG backend. 'mcp:server:tool' routes to an MCP server "
|
||||
"(e.g. 'mcp:search:web_search').",
|
||||
),
|
||||
SettingDef(
|
||||
"tools.rerank_url",
|
||||
"str",
|
||||
"",
|
||||
"Full rerank endpoint URL (empty = reranking disabled)",
|
||||
"tools",
|
||||
help="Full URL of a Cohere/Jina-compatible rerank endpoint, including the path "
|
||||
"(e.g. 'http://vllm:8000/rerank', 'http://tei:8080/rerank', or "
|
||||
"'https://api.cohere.com/v2/rerank'). Turnstone runs no rerank model itself: "
|
||||
"web_search results and large web_fetch pages are reranked by POSTing to this "
|
||||
"endpoint. Empty disables reranking entirely. Overrides config.toml [tools] "
|
||||
"rerank_url and $TURNSTONE_RERANK_URL.",
|
||||
),
|
||||
SettingDef(
|
||||
"tools.rerank_model",
|
||||
"str",
|
||||
"",
|
||||
"Reranker model name sent to the endpoint (empty = endpoint default)",
|
||||
"tools",
|
||||
help="Model identifier passed in the rerank request body (e.g. "
|
||||
"'BAAI/bge-reranker-v2-m3', 'rerank-v3.5'). Empty uses the endpoint's default "
|
||||
"model. Overrides config.toml [tools] rerank_model.",
|
||||
),
|
||||
SettingDef(
|
||||
"tools.rerank_api_key",
|
||||
"str",
|
||||
"",
|
||||
"Bearer token for the rerank endpoint (write-only)",
|
||||
"tools",
|
||||
is_secret=True,
|
||||
help="Sent as 'Authorization: Bearer <key>' to the rerank endpoint. Required for "
|
||||
"hosted providers (Cohere/Jina/Voyage); usually empty for a self-hosted vLLM/TEI "
|
||||
"instance. Write-only: never returned by the API. For non-console deployments set "
|
||||
"it in config.toml [tools] rerank_api_key (secrets belong in config.toml, not env).",
|
||||
),
|
||||
SettingDef(
|
||||
"tools.rerank_web_search",
|
||||
"bool",
|
||||
True,
|
||||
"Rerank web_search results when a rerank endpoint is configured",
|
||||
"tools",
|
||||
help="When a rerank endpoint is set (tools.rerank_url), re-order web_search "
|
||||
"results by query relevance before returning the top hits. No effect when no "
|
||||
"endpoint is configured. Disable to keep the search backend's native ranking.",
|
||||
),
|
||||
SettingDef(
|
||||
"tools.reranker_alias",
|
||||
"str",
|
||||
"",
|
||||
"Model alias of the reranker (a model with supports_rerank), or empty",
|
||||
"tools",
|
||||
help="Selects a reranker added in the admin Models tab: create a model definition "
|
||||
'with capability {"supports_rerank": true} and its base_url set to a Cohere/Jina-'
|
||||
"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).",
|
||||
),
|
||||
# -- server ---------------------------------------------------------
|
||||
SettingDef(
|
||||
"server.workstream_idle_timeout",
|
||||
|
||||
@@ -16,6 +16,7 @@ clients.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, Protocol
|
||||
|
||||
import httpx
|
||||
@@ -27,6 +28,15 @@ if TYPE_CHECKING:
|
||||
|
||||
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.
|
||||
_RERANK_POOL = 50
|
||||
|
||||
|
||||
class WebSearchClient(Protocol):
|
||||
"""Minimal interface for a web search backend."""
|
||||
@@ -67,7 +77,7 @@ class SearXNGClient:
|
||||
timeout=self._timeout,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return _format_searxng(resp.json(), query, max_results)
|
||||
return _format_searxng(resp.json(), query, max_results, reranker=kwargs.get("reranker"))
|
||||
|
||||
|
||||
class MCPSearchClient:
|
||||
@@ -95,7 +105,12 @@ class MCPSearchClient:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _format_searxng(data: dict[str, Any], query: str, max_results: int = 5) -> str:
|
||||
def _format_searxng(
|
||||
data: dict[str, Any],
|
||||
query: str,
|
||||
max_results: int = 5,
|
||||
reranker: Reranker | None = None,
|
||||
) -> str:
|
||||
parts: list[str] = []
|
||||
|
||||
# Instant answers (calculator, Wikipedia summaries, …) — engine-dependent,
|
||||
@@ -119,6 +134,8 @@ def _format_searxng(data: dict[str, Any], query: str, max_results: int = 5) -> s
|
||||
|
||||
results = data.get("results") or []
|
||||
if results:
|
||||
if reranker is not None and len(results) > 1:
|
||||
results = _rerank_results(query, results, reranker)
|
||||
lines = []
|
||||
for i, r in enumerate(results[:max_results], 1):
|
||||
title = r.get("title", "")
|
||||
@@ -142,6 +159,38 @@ def _format_searxng(data: dict[str, Any], query: str, max_results: int = 5) -> s
|
||||
return f"No results for '{query}'."
|
||||
|
||||
|
||||
def _rerank_results(
|
||||
query: str, results: list[dict[str, Any]], reranker: Reranker
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Reorder SearxNG ``results`` by query relevance using ``reranker``.
|
||||
|
||||
Sends at most ``_RERANK_POOL`` results (title + snippet) to the reranker and
|
||||
splices its ordering back in. Falls back to the original SearxNG order on any
|
||||
error or if the reranker returns nothing usable — reranking must never make
|
||||
web_search fail or silently drop results.
|
||||
"""
|
||||
pool = results[:_RERANK_POOL]
|
||||
tail = results[_RERANK_POOL:]
|
||||
try:
|
||||
docs = [f"{r.get('title', '')}\n{r.get('content') or ''}".strip() for r in pool]
|
||||
order = reranker(query, docs)
|
||||
except Exception as e:
|
||||
log.warning("rerank failed; using native result order: %s", e)
|
||||
return results
|
||||
seen: set[int] = set()
|
||||
reordered: list[dict[str, Any]] = []
|
||||
for idx in order:
|
||||
if isinstance(idx, int) and 0 <= idx < len(pool) and idx not in seen:
|
||||
seen.add(idx)
|
||||
reordered.append(pool[idx])
|
||||
if not reordered:
|
||||
return results
|
||||
# Keep any pool items the reranker omitted (e.g. a top_n subset), then the
|
||||
# un-reranked tail beyond the pool cap.
|
||||
reordered.extend(pool[i] for i in range(len(pool)) if i not in seen)
|
||||
return reordered + tail
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Resolver
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user