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:
@@ -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__}")
|
||||
|
||||
Reference in New Issue
Block a user