feat(notion-search): add Claude Haiku rerank module
Builds a rerank prompt with title + flattened properties + excerpt for each candidate, calls Claude, parses JSON, sorts by score descending, takes top N. On any failure (LLM error, missing JSON, parse error, non-list shape), falls back to candidates in input order with null relevance/snippet. 4 tests: ordering, limit, parse-error fallback, exception fallback. 26 passing total. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -236,3 +236,97 @@ def enrich_candidates(notion, candidates: List[Dict]) -> List[Dict]:
|
||||
"excerpt": excerpt,
|
||||
})
|
||||
return enriched
|
||||
|
||||
|
||||
RERANK_PROMPT = """You are a reranker for Notion semantic search. Score each candidate 0.0-1.0 by how relevant it is to the user's ORIGINAL query (not any expanded variants).
|
||||
|
||||
User's query: {query}
|
||||
|
||||
Candidates:
|
||||
{candidates}
|
||||
|
||||
Return ONLY a JSON array, ordered however you like. Each object has:
|
||||
- "index": integer matching the candidate's [N] number
|
||||
- "score": float 0.0-1.0
|
||||
- "why": one short sentence (under 80 chars) explaining the score
|
||||
|
||||
Example output:
|
||||
[{{"index": 0, "score": 0.95, "why": "Direct match — covers exactly this topic"}},
|
||||
{{"index": 2, "score": 0.6, "why": "Adjacent — shares context but not topic"}}]"""
|
||||
|
||||
|
||||
def _format_candidate_for_rerank(idx: int, c: Dict) -> str:
|
||||
parts = [f"[{idx}] {c['title']}"]
|
||||
props_str = ", ".join(f"{k}: {v}" for k, v in c.get("properties", {}).items())
|
||||
if props_str:
|
||||
parts.append(f" Properties: {props_str}")
|
||||
if c.get("excerpt"):
|
||||
parts.append(f" Excerpt: {c['excerpt']}")
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
def _fallback_rank(candidates: List[Dict], limit: int) -> List[Dict]:
|
||||
"""Return candidates in input order with null relevance/snippet."""
|
||||
return [
|
||||
{**c, "relevance": None, "snippet": None}
|
||||
for c in candidates[:limit]
|
||||
]
|
||||
|
||||
|
||||
def rerank(
|
||||
query: str,
|
||||
candidates: List[Dict],
|
||||
*,
|
||||
llm_caller: Optional[Callable[..., str]] = None,
|
||||
limit: int = 10,
|
||||
) -> List[Dict]:
|
||||
"""Rerank candidates against the original query. Fallback to input order on failure."""
|
||||
if llm_caller is None:
|
||||
llm_caller = default_llm_caller
|
||||
if not candidates:
|
||||
return []
|
||||
|
||||
formatted = "\n\n".join(
|
||||
_format_candidate_for_rerank(i, c) for i, c in enumerate(candidates)
|
||||
)
|
||||
prompt = RERANK_PROMPT.format(query=query, candidates=formatted)
|
||||
|
||||
try:
|
||||
response = llm_caller(prompt)
|
||||
except Exception as exc:
|
||||
print(f"Warning: rerank failed ({exc}); returning unranked results",
|
||||
file=sys.stderr)
|
||||
return _fallback_rank(candidates, limit)
|
||||
|
||||
match = re.search(r'\[.*\]', response, re.DOTALL)
|
||||
if not match:
|
||||
print("Warning: rerank returned no JSON array; returning unranked results",
|
||||
file=sys.stderr)
|
||||
return _fallback_rank(candidates, limit)
|
||||
|
||||
try:
|
||||
scored = json.loads(match.group(0))
|
||||
except json.JSONDecodeError:
|
||||
print("Warning: rerank returned invalid JSON; returning unranked results",
|
||||
file=sys.stderr)
|
||||
return _fallback_rank(candidates, limit)
|
||||
|
||||
if not isinstance(scored, list):
|
||||
print("Warning: rerank returned non-list shape; returning unranked results",
|
||||
file=sys.stderr)
|
||||
return _fallback_rank(candidates, limit)
|
||||
|
||||
# Sort by score descending and map back to candidates
|
||||
scored.sort(key=lambda s: s.get("score", 0), reverse=True)
|
||||
out = []
|
||||
for entry in scored[:limit]:
|
||||
idx = entry.get("index")
|
||||
if not isinstance(idx, int) or idx < 0 or idx >= len(candidates):
|
||||
continue
|
||||
c = candidates[idx]
|
||||
out.append({
|
||||
**c,
|
||||
"relevance": float(entry.get("score", 0.0)),
|
||||
"snippet": entry.get("why", ""),
|
||||
})
|
||||
return out
|
||||
|
||||
Reference in New Issue
Block a user