From c72931b30fa8cd6f27e2c061187a00cc21011670 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Thu, 17 Sep 2026 15:52:34 +0800 Subject: [PATCH] feat(python): add TypeSafe reranker (#4209) Adds `TypeSafeReranker`, which reranks vector, FTS, and hybrid results with the [TypeSafe System One API](https://docs.typesafe.ai/introduction). Each result is scored independently: TypeSafe reads `{"query", "document"}` and answers a yes/no (noul) question, and the probability of yes becomes `_relevance_score`. Unlike listwise LLM rerankers, the score is an absolute probability, so it is comparable across queries and can be thresholded. The question's `instructions` and `true`/`false` `criteria` are configurable, since domain-specific criteria are what make this kind of scoring work well ([TypeSafe's re-ranking cookbook](https://docs.typesafe.ai/cookbooks/rerank_typesafe)). The API takes one state per request, so the reranker sends one request per result on a thread pool bounded by `max_concurrency`. It deliberately does not use the background event loop: rerankers are called synchronously from inside the async query APIs, where `LOOP.run` would deadlock. TypeSafe scores for the same pair vary slightly between calls, so results with close scores can swap places when a search is repeated. The shared reranker test helper now takes `deterministic=False` for this case: it still checks result sizes and descending scores, but not that two identical searches return the same order. The SDK is imported only when the client is created and questions are sent as plain dicts, so the new tests run in CI with a fake client and without `typesafe-sdk` installed. The live-API test is skipped without `TYPESAFE_API_KEY`. Ranking quality has not been compared with other API rerankers. --- python/pyproject.toml | 1 + python/python/lancedb/rerankers/__init__.py | 2 + python/python/lancedb/rerankers/typesafe.py | 159 ++++++++++++++++++++ python/python/tests/test_rerankers.py | 107 ++++++++++++- 4 files changed, 266 insertions(+), 3 deletions(-) create mode 100644 python/python/lancedb/rerankers/typesafe.py diff --git a/python/pyproject.toml b/python/pyproject.toml index dc5ea603d..de84367a8 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -126,6 +126,7 @@ include = [ "python/lancedb/rerankers/util.py", "python/lancedb/rerankers/__init__.py", "python/lancedb/rerankers/voyageai.py", + "python/lancedb/rerankers/typesafe.py", "python/lancedb/rerankers/jinaai.py", "python/lancedb/rerankers/openai.py", "python/lancedb/rerankers/cross_encoder.py", diff --git a/python/python/lancedb/rerankers/__init__.py b/python/python/lancedb/rerankers/__init__.py index aef2d331e..bd906f831 100644 --- a/python/python/lancedb/rerankers/__init__.py +++ b/python/python/lancedb/rerankers/__init__.py @@ -11,6 +11,7 @@ from .jinaai import JinaReranker from .rrf import RRFReranker from .mrr import MRRReranker from .answerdotai import AnswerdotaiRerankers +from .typesafe import TypeSafeReranker from .voyageai import VoyageAIReranker from .watsonx import WatsonxReranker @@ -27,6 +28,7 @@ __all__ = [ "JinaReranker", "RRFReranker", "AnswerdotaiRerankers", + "TypeSafeReranker", "VoyageAIReranker", "MRRReranker", "WatsonxReranker", diff --git a/python/python/lancedb/rerankers/typesafe.py b/python/python/lancedb/rerankers/typesafe.py new file mode 100644 index 000000000..2da9b673c --- /dev/null +++ b/python/python/lancedb/rerankers/typesafe.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The LanceDB Authors + + +from concurrent.futures import ThreadPoolExecutor +from functools import cached_property +from typing import Any, Dict, Mapping, Optional + +import pyarrow as pa + +from ..util import attempt_import_or_raise +from .base import Reranker + +DEFAULT_INSTRUCTIONS = ( + "Does the document contain information that answers or directly addresses " + "the query?" +) +DEFAULT_CRITERIA = { + "true": "The document answers the query or states facts the query asks about.", + "false": "The document is off-topic, or only shares keywords or a general " + "subject with the query without addressing it.", +} +# The key TypeSafe returns the answer under. It is not sent to the model. +_QUESTION_ID = "relevance" + + +class TypeSafeReranker(Reranker): + """ + Reranks the results using the TypeSafe System One API. + https://docs.typesafe.ai/cookbooks/rerank_typesafe + + Each result is scored independently: TypeSafe reads the query together with + the result's ``column`` value and answers a yes/no ("noul") question about + whether the document is relevant. The probability of "yes", between 0 and 1, + becomes the ``_relevance_score``. Because every score is an absolute + probability rather than a position in the list, scores are comparable + across queries and can be thresholded. They are model estimates, though, + and can vary slightly between identical calls, so results with close scores + may swap places when the same search is repeated. + + One request is sent per result, up to ``max_concurrency`` at a time. + + Parameters + ---------- + model_name : str, default "jev-latest" + The TypeSafe model to use. + column : str, default "text" + The name of the column holding the document text to score. + instructions : str, optional + The yes/no question asked about each query and document pair. The state + TypeSafe reads is ``{"query": , "document": }``. + Defaults to a generic relevance question. + criteria : Mapping[str, str], optional + What a "yes" and a "no" mean, as a mapping with the keys ``"true"`` and + ``"false"``. Domain-specific criteria usually rank better than the + generic default. Pass an empty mapping to send no criteria. + return_score : str, default "relevance" + Options are "relevance" or "all". If "all", keeps the vector and FTS + scores alongside the relevance score. + api_key : str, optional + The API key to use. If None, the TypeSafe SDK reads the + ``TYPESAFE_API_KEY`` environment variable. + max_concurrency : int, default 8 + The maximum number of TypeSafe requests in flight for one rerank call. + """ + + def __init__( + self, + model_name: str = "jev-latest", + column: str = "text", + instructions: Optional[str] = None, + criteria: Optional[Mapping[str, str]] = None, + return_score: str = "relevance", + api_key: Optional[str] = None, + max_concurrency: int = 8, + ): + super().__init__(return_score) + if max_concurrency < 1: + raise ValueError("max_concurrency must be at least 1") + criteria = DEFAULT_CRITERIA if criteria is None else dict(criteria) + unknown = set(criteria) - {"true", "false"} + if unknown: + raise ValueError( + f"criteria keys must be 'true' or 'false', got {sorted(unknown)}" + ) + self.model_name = model_name + self.column = column + self.instructions = instructions or DEFAULT_INSTRUCTIONS + self.criteria = criteria + self.api_key = api_key + self.max_concurrency = max_concurrency + + def __str__(self): + return f"TypeSafeReranker(model_name={self.model_name})" + + @cached_property + def _client(self): + typesafe_sdk = attempt_import_or_raise("typesafe_sdk", "typesafe-sdk") + return typesafe_sdk.TypeSafeClient(api_key=self.api_key) + + @cached_property + def _question(self) -> Dict[str, Any]: + question: Dict[str, Any] = {"type": "noul", "instructions": self.instructions} + if self.criteria: + question["criteria"] = self.criteria + return question + + def _score(self, query: str, document: Optional[str]) -> float: + if document is None: + return 0.0 + response = self._client.system_one( + state={"query": query, "document": document}, + questions={_QUESTION_ID: self._question}, + model=self.model_name, + ) + return response.answers[_QUESTION_ID].noul + + def _rerank(self, result_set: pa.Table, query: str) -> pa.Table: + result_set = self._handle_empty_results(result_set) + if len(result_set) == 0: + return result_set + docs = result_set[self.column].to_pylist() + # Rerankers are also called synchronously from inside the async query + # APIs, so the requests run on threads rather than on an event loop. + with ThreadPoolExecutor( + max_workers=min(self.max_concurrency, len(docs)) + ) as pool: + scores = list(pool.map(lambda doc: self._score(query, doc), docs)) + result_set = result_set.append_column( + "_relevance_score", pa.array(scores, type=pa.float32()) + ) + return result_set.sort_by([("_relevance_score", "descending")]) + + def rerank_hybrid( + self, + query: str, + vector_results: pa.Table, + fts_results: pa.Table, + ): + if self.score == "all": + combined_results = self._merge_and_keep_scores(vector_results, fts_results) + else: + combined_results = self.merge_results(vector_results, fts_results) + combined_results = self._rerank(combined_results, query) + if self.score == "relevance": + combined_results = self._keep_relevance_score(combined_results) + return combined_results + + def rerank_vector(self, query: str, vector_results: pa.Table): + vector_results = self._rerank(vector_results, query) + if self.score == "relevance": + vector_results = vector_results.drop_columns(["_distance"]) + return vector_results + + def rerank_fts(self, query: str, fts_results: pa.Table): + fts_results = self._rerank(fts_results, query) + if self.score == "relevance": + fts_results = fts_results.drop_columns(["_score"]) + return fts_results diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index 4430fa98f..290164a13 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -21,6 +21,7 @@ from lancedb.rerankers import ( OpenaiReranker, JinaReranker, AnswerdotaiRerankers, + TypeSafeReranker, VoyageAIReranker, MRRReranker, WatsonxReranker, @@ -101,7 +102,15 @@ def get_test_table(tmp_path): return table, MyTable -def _run_test_reranker(reranker, table, query, query_vector, schema): +def _run_test_reranker( + reranker, table, query, query_vector, schema, deterministic=True +): + """Exercise a reranker across search types. + + Set ``deterministic=False`` for rerankers whose scores can vary between + identical calls, such as remote model APIs; repeated searches are then not + expected to return the same order. + """ # Hybrid search setting result1 = ( table.search(query, query_type="hybrid", vector_column_name="vector") @@ -113,7 +122,8 @@ def _run_test_reranker(reranker, table, query, query_vector, schema): .rerank(reranker=reranker) .to_pydantic(schema) ) - assert result1 == result2 + if deterministic: + assert result1 == result2 query_vector = table.to_pandas()["vector"][0] result = ( @@ -203,7 +213,9 @@ def _run_test_reranker(reranker, table, query, query_vector, schema): ) assert len(result_deduped) <= 20 result_arrow = reranker.rerank_multivector([rs1.to_arrow(), rs2.to_arrow()], query) - assert len(result) == 20 and result == result_arrow + assert len(result) == 20 and len(result_arrow) == 20 + if deterministic: + assert result == result_arrow def _run_test_hybrid_reranker(reranker, tmp_path): @@ -515,6 +527,95 @@ def test_voyageai_reranker(tmp_path): _run_test_reranker(reranker, table, "single player experience", None, schema) +class _FakeTypeSafeClient: + """Stands in for ``typesafe_sdk.TypeSafeClient``; scores by word overlap.""" + + def __init__(self): + self.requests = [] + + def system_one(self, state, questions, model): + self.requests.append((state, questions, model)) + query_words = set(state["query"].lower().split()) + doc_words = set(state["document"].lower().split()) + noul = len(query_words & doc_words) / len(query_words) + answers = {key: type("NoulAnswer", (), {"noul": noul})() for key in questions} + return type("SystemOneResponse", (), {"answers": answers})() + + +def test_typesafe_reranker_with_fake_client(tmp_path): + reranker = TypeSafeReranker(max_concurrency=4) + reranker._client = _FakeTypeSafeClient() + table, schema = get_test_table(tmp_path) + _run_test_reranker(reranker, table, "single player experience", None, schema) + + state, questions, model = reranker._client.requests[0] + assert model == "jev-latest" + assert set(state) == {"query", "document"} + assert questions == { + "relevance": { + "type": "noul", + "instructions": reranker.instructions, + "criteria": reranker.criteria, + } + } + + +def test_typesafe_reranker_scores_each_row(): + reranker = TypeSafeReranker( + column="body", + instructions="Is this about cats?", + criteria={}, + return_score="all", + ) + reranker._client = _FakeTypeSafeClient() + results = pa.table( + { + "body": ["dogs bark", "cats purr and cats nap", None, "cats"], + "_distance": [0.1, 0.2, 0.3, 0.4], + } + ) + + reranked = reranker.rerank_vector("cats nap", results) + + assert reranked["body"].to_pylist() == [ + "cats purr and cats nap", + "cats", + "dogs bark", + None, + ] + assert reranked["_relevance_score"].to_pylist() == [1.0, 0.5, 0.0, 0.0] + assert reranked["_distance"].to_pylist() == [0.2, 0.4, 0.1, 0.3] + # Null documents are scored 0 without a request. + assert len(reranker._client.requests) == 3 + assert reranker._client.requests[0][1] == { + "relevance": {"type": "noul", "instructions": "Is this about cats?"} + } + + +def test_typesafe_reranker_rejects_invalid_arguments(): + with pytest.raises(ValueError, match="criteria keys"): + TypeSafeReranker(criteria={"yes": "relevant"}) + with pytest.raises(ValueError, match="max_concurrency"): + TypeSafeReranker(max_concurrency=0) + + +@pytest.mark.skipif( + os.environ.get("TYPESAFE_API_KEY") is None, reason="TYPESAFE_API_KEY not set" +) +def test_typesafe_reranker(tmp_path): + pytest.importorskip("typesafe_sdk") + reranker = TypeSafeReranker() + table, schema = get_test_table(tmp_path) + _run_test_reranker( + reranker, + table, + "single player experience", + None, + schema, + deterministic=False, + ) + + def test_empty_result_reranker(): pytest.importorskip("sentence_transformers") db = lancedb.connect("memory://")