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:
2026-04-28 08:01:45 +09:00
parent e67209b905
commit 20029ecc9c
2 changed files with 172 additions and 0 deletions

View File

@@ -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

View File

@@ -383,6 +383,80 @@ def test_enrich_keeps_falsy_but_meaningful_values():
"number=0 kept (not filtered as 'empty')")
def test_rerank_orders_by_score_descending():
"""LLM returns scores; output sorted score-desc, top N."""
import notion_search
candidates = [
{"id": "p1", "url": "u1", "title": "Page 1", "properties": {}, "excerpt": "x"},
{"id": "p2", "url": "u2", "title": "Page 2", "properties": {}, "excerpt": "x"},
{"id": "p3", "url": "u3", "title": "Page 3", "properties": {}, "excerpt": "x"},
]
# LLM returns: p3 most relevant, then p1, then p2
fake_llm = lambda prompt, **kw: '''[
{"index": 2, "score": 0.95, "why": "best match"},
{"index": 0, "score": 0.7, "why": "ok match"},
{"index": 1, "score": 0.3, "why": "weak match"}
]'''
ranked = notion_search.rerank("query", candidates, llm_caller=fake_llm, limit=3)
_assert(ranked[0]["id"] == "p3", "highest score first")
_assert(ranked[0]["relevance"] == 0.95, "score attached to result")
_assert(ranked[0]["snippet"] == "best match", "why text attached as snippet")
_assert(ranked[1]["id"] == "p1", "second place correct")
_assert(ranked[2]["id"] == "p2", "lowest score last")
def test_rerank_respects_limit():
"""limit=2 → returns top 2 only."""
import notion_search
candidates = [
{"id": f"p{i}", "url": f"u{i}", "title": f"Page {i}", "properties": {}, "excerpt": "x"}
for i in range(5)
]
fake_llm = lambda prompt, **kw: '''[
{"index": 0, "score": 0.9, "why": "x"},
{"index": 1, "score": 0.8, "why": "x"},
{"index": 2, "score": 0.7, "why": "x"},
{"index": 3, "score": 0.6, "why": "x"},
{"index": 4, "score": 0.5, "why": "x"}
]'''
ranked = notion_search.rerank("query", candidates, llm_caller=fake_llm, limit=2)
_assert(len(ranked) == 2, "exactly limit results returned")
def test_rerank_falls_back_on_parse_error():
"""Non-JSON response → return candidates in input order, unranked (no scores)."""
import notion_search
candidates = [
{"id": "p1", "url": "u1", "title": "Page 1", "properties": {}, "excerpt": "x"},
{"id": "p2", "url": "u2", "title": "Page 2", "properties": {}, "excerpt": "x"},
]
fake_llm = lambda prompt, **kw: "I cannot help with that."
ranked = notion_search.rerank("query", candidates, llm_caller=fake_llm, limit=10)
_assert(len(ranked) == 2, "all candidates returned in fallback")
_assert(ranked[0]["id"] == "p1", "fallback preserves input order")
_assert(ranked[0]["relevance"] is None, "no score in fallback")
_assert(ranked[0]["snippet"] is None, "no snippet in fallback")
def test_rerank_falls_back_on_llm_exception():
"""LLM raises → fallback unranked."""
import notion_search
candidates = [{"id": "p1", "url": "u1", "title": "Page 1", "properties": {}, "excerpt": "x"}]
def boom(prompt, **kw):
raise RuntimeError("API down")
ranked = notion_search.rerank("query", candidates, llm_caller=boom, limit=10)
_assert(len(ranked) == 1, "candidate returned despite LLM failure")
_assert(ranked[0]["relevance"] is None, "no score in fallback")
def run_all():
tests = [
test_call_claude_dispatches_to_sdk_when_available,
@@ -407,6 +481,10 @@ def run_all():
test_enrich_falls_back_to_empty_excerpt,
test_enrich_truncates_long_excerpt,
test_enrich_keeps_falsy_but_meaningful_values,
test_rerank_orders_by_score_descending,
test_rerank_respects_limit,
test_rerank_falls_back_on_parse_error,
test_rerank_falls_back_on_llm_exception,
]
for t in tests:
print(f"\n{t.__name__}")