chaim
bb9934dde4
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>
114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
"""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
|
|
from api.services.kb.db import get_pool
|
|
|
|
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(
|
|
query: str,
|
|
kind: Literal["law", "regulation", "circular", "any"] = "any",
|
|
top_k: int = 8,
|
|
) -> list[dict]:
|
|
query = (query or "").strip()
|
|
if not query:
|
|
return []
|
|
|
|
[vec] = await voyage.embed([query], input_type="query")
|
|
|
|
kind_filter = ""
|
|
params: list = [vec, query, _CANDIDATES_PER_SIDE]
|
|
if kind != "any":
|
|
kind_filter = "AND s.kind = $4"
|
|
params.append(kind)
|
|
|
|
# Two CTEs — vector rank and lexical rank — fused via RRF.
|
|
sql = f"""
|
|
WITH vec AS (
|
|
SELECT c.id, ROW_NUMBER() OVER (ORDER BY c.embedding <=> $1::vector) AS rank
|
|
FROM kb_chunk c
|
|
JOIN kb_source s ON s.id = c.source_id
|
|
WHERE c.embedding IS NOT NULL
|
|
AND s.superseded_by IS NULL
|
|
{kind_filter}
|
|
ORDER BY c.embedding <=> $1::vector
|
|
LIMIT $3
|
|
),
|
|
lex AS (
|
|
SELECT c.id, ROW_NUMBER() OVER (ORDER BY ts_rank_cd(c.content_tsv, q) DESC) AS rank
|
|
FROM kb_chunk c
|
|
JOIN kb_source s ON s.id = c.source_id,
|
|
websearch_to_tsquery('simple', $2) AS q
|
|
WHERE c.content_tsv @@ q
|
|
AND s.superseded_by IS NULL
|
|
{kind_filter}
|
|
ORDER BY ts_rank_cd(c.content_tsv, q) DESC
|
|
LIMIT $3
|
|
),
|
|
fused AS (
|
|
SELECT id, SUM(score) AS score FROM (
|
|
SELECT id, 1.0 / ({_RRF_K} + rank) AS score FROM vec
|
|
UNION ALL
|
|
SELECT id, 1.0 / ({_RRF_K} + rank) AS score FROM lex
|
|
) u GROUP BY id
|
|
)
|
|
SELECT
|
|
s.kind, s.title, s.identifier, s.source_url,
|
|
s.published_at, s.effective_at,
|
|
c.heading_path, c.section_ref, c.content,
|
|
f.score
|
|
FROM fused f
|
|
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 {_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 RRF_candidates=%d",
|
|
query[:80], kind, len(candidates),
|
|
)
|
|
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
|