diff --git a/docs/tools.md b/docs/tools.md index 2ea2bbd2..b0ed131a 100644 --- a/docs/tools.md +++ b/docs/tools.md @@ -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 diff --git a/tests/test_rerank.py b/tests/test_rerank.py new file mode 100644 index 00000000..01615eb5 --- /dev/null +++ b/tests/test_rerank.py @@ -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 diff --git a/tests/test_web_search.py b/tests/test_web_search.py index 202f21f0..b0c3379f 100644 --- a/tests/test_web_search.py +++ b/tests/test_web_search.py @@ -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 diff --git a/turnstone.example.toml b/turnstone.example.toml index 7e84f63c..e920004f 100644 --- a/turnstone.example.toml +++ b/turnstone.example.toml @@ -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) --- diff --git a/turnstone/console/static/admin.js b/turnstone/console/static/admin.js index f4f909ae..597e732e 100644 --- a/turnstone/console/static/admin.js +++ b/turnstone/console/static/admin.js @@ -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 diff --git a/turnstone/core/config.py b/turnstone/core/config.py index 5e14538c..dd42ce27 100644 --- a/turnstone/core/config.py +++ b/turnstone/core/config.py @@ -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) diff --git a/turnstone/core/rerank.py b/turnstone/core/rerank.py new file mode 100644 index 00000000..990b158a --- /dev/null +++ b/turnstone/core/rerank.py @@ -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": "", "query": "", "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 `` 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, + ) diff --git a/turnstone/core/session.py b/turnstone/core/session.py index 1b4daed2..8c089ecb 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -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) diff --git a/turnstone/core/settings_registry.py b/turnstone/core/settings_registry.py index ecb703bd..3d9d4525 100644 --- a/turnstone/core/settings_registry.py +++ b/turnstone/core/settings_registry.py @@ -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 ' 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", diff --git a/turnstone/core/web_search.py b/turnstone/core/web_search.py index adfa07a3..05e21650 100644 --- a/turnstone/core/web_search.py +++ b/turnstone/core/web_search.py @@ -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 # ---------------------------------------------------------------------------