feat(kb): Voyage rerank-2.5 cross-encoder after RRF

Adds a rerank stage between the hybrid RRF fusion and the final top_k cut.
RRF is a heuristic over ranks — it can't tell which of two equally-ranked
candidates is actually more relevant. A cross-encoder scores each
(query, passage) pair directly and produces a true relevance order.

- voyage.rerank(query, documents, top_k) → list of {index, relevance_score}
  sorted high→low. Defaults to the model in VOYAGE_RERANK_MODEL
  (rerank-2.5, 32K tokens per pair, multilingual incl. Hebrew).
- search pulls top-30 from RRF, passes them to rerank, then returns the
  reranker's top-k. Each returned item gets a rerank_score field.
- Fallback: if the rerank API errors (auth, 5xx, timeout), fall back to
  the RRF ordering and log a warning — degraded precision beats outage.
- KB_RERANK_DISABLED=1 bypasses rerank for debugging.

Refs Task Master #2

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-04-21 16:45:02 +00:00
parent 68c9506f9c
commit bb9934dde4
2 changed files with 80 additions and 5 deletions
+34 -5
View File
@@ -1,8 +1,10 @@
"""Hybrid search: vector + full-text, fused with Reciprocal Rank Fusion."""
"""Hybrid search: vector + full-text, fused with Reciprocal Rank Fusion,
and finally reranked with a Voyage cross-encoder for precision."""
from __future__ import annotations
import logging
import os
from typing import Literal
from api.services.kb import voyage
@@ -12,6 +14,9 @@ logger = logging.getLogger("shira.kb.search")
_RRF_K = 60 # standard RRF constant
_CANDIDATES_PER_SIDE = 30
# Number of fused RRF candidates fed to the reranker. Higher = better
# recall before the final cut; capped at ~30 to keep latency reasonable.
_RERANK_POOL = 30
async def search(
@@ -70,15 +75,39 @@ async def search(
JOIN kb_chunk c ON c.id = f.id
JOIN kb_source s ON s.id = c.source_id
ORDER BY f.score DESC
LIMIT {int(top_k)}
LIMIT {_RERANK_POOL}
"""
pool = await get_pool()
async with pool.acquire() as conn:
rows = await conn.fetch(sql, *params)
candidates = [dict(r) for r in rows]
logger.info(
"[kb.search] query=%r kind=%s hits=%d",
query[:80], kind, len(rows),
"[kb.search] query=%r kind=%s RRF_candidates=%d",
query[:80], kind, len(candidates),
)
return [dict(r) for r in rows]
if not candidates:
return []
rerank_disabled = os.environ.get("KB_RERANK_DISABLED", "").lower() in ("1", "true")
if rerank_disabled or len(candidates) <= top_k:
return candidates[:top_k]
# Cross-encoder rerank for precision. On failure fall back to the RRF order
# rather than returning nothing — degraded precision is better than outage.
texts = [c["content"] for c in candidates]
try:
ranked = await voyage.rerank(query=query, documents=texts, top_k=top_k)
except voyage.VoyageError as e:
logger.warning("[kb.search] rerank failed, falling back to RRF: %s", e)
return candidates[:top_k]
out: list[dict] = []
for r in ranked:
i = r["index"]
if 0 <= i < len(candidates):
item = {**candidates[i], "rerank_score": r["relevance_score"]}
out.append(item)
logger.info("[kb.search] rerank kept top %d / %d", len(out), len(candidates))
return out
+46
View File
@@ -11,6 +11,7 @@ import httpx
logger = logging.getLogger("shira.kb.voyage")
_VOYAGE_URL = "https://api.voyageai.com/v1/embeddings"
_VOYAGE_RERANK_URL = "https://api.voyageai.com/v1/rerank"
_MAX_BATCH = 128
# Voyage caps each request at 120K tokens total. Keep a safety margin so our
# char-based estimate doesn't push over when the real tokenizer disagrees.
@@ -70,3 +71,48 @@ async def embed(
bi + 1, len(batches), len(batch), input_type,
)
return out
async def rerank(
query: str,
documents: list[str],
top_k: int | None = None,
model: str | None = None,
) -> list[dict]:
"""Score each document against the query using a Voyage reranker.
Returns a list of {index, relevance_score} sorted high→low. `index`
refers back into the original `documents` list so the caller can map
scores back to its own metadata.
Empty input → empty output. Upstream failures raise VoyageError; callers
that want a soft fallback should catch it.
"""
if not query or not documents:
return []
api_key = os.environ.get("VOYAGE_API_KEY")
if not api_key:
raise VoyageError("VOYAGE_API_KEY is not set")
model = model or os.environ.get("VOYAGE_RERANK_MODEL", "rerank-2.5")
payload: dict = {
"model": model,
"query": query,
"documents": documents,
}
if top_k is not None:
payload["top_k"] = int(top_k)
async with httpx.AsyncClient(timeout=60) as http:
resp = await http.post(
_VOYAGE_RERANK_URL,
headers={"Authorization": f"Bearer {api_key}"},
json=payload,
)
if resp.status_code != 200:
raise VoyageError(f"Voyage rerank {resp.status_code}: {resp.text[:200]}")
data = resp.json().get("data", [])
return [
{"index": item["index"], "relevance_score": item["relevance_score"]}
for item in data
]