From bb9934dde4c6c6bdfd8da7abaebbacf1dc10c27e Mon Sep 17 00:00:00 2001 From: Chaim Date: Tue, 21 Apr 2026 16:45:02 +0000 Subject: [PATCH] feat(kb): Voyage rerank-2.5 cross-encoder after RRF MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- api/services/kb/search.py | 39 ++++++++++++++++++++++++++++----- api/services/kb/voyage.py | 46 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 5 deletions(-) diff --git a/api/services/kb/search.py b/api/services/kb/search.py index c456ae1..1549ab1 100644 --- a/api/services/kb/search.py +++ b/api/services/kb/search.py @@ -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 diff --git a/api/services/kb/voyage.py b/api/services/kb/voyage.py index c492e86..0f4542c 100644 --- a/api/services/kb/voyage.py +++ b/api/services/kb/voyage.py @@ -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 + ]