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

@@ -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__}")