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.
This commit is contained in:
Xuanwo
2026-09-17 15:52:34 +08:00
committed by GitHub
parent 60a1b4c219
commit c72931b30f
4 changed files with 266 additions and 3 deletions
+104 -3
View File
@@ -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://")