diff --git a/custom-skills/31-notion-organizer/code/scripts/notion_search.py b/custom-skills/31-notion-organizer/code/scripts/notion_search.py index 8098fd8..46f7db1 100644 --- a/custom-skills/31-notion-organizer/code/scripts/notion_search.py +++ b/custom-skills/31-notion-organizer/code/scripts/notion_search.py @@ -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 diff --git a/custom-skills/31-notion-organizer/code/scripts/test_notion_search.py b/custom-skills/31-notion-organizer/code/scripts/test_notion_search.py index 02206e9..45464b9 100644 --- a/custom-skills/31-notion-organizer/code/scripts/test_notion_search.py +++ b/custom-skills/31-notion-organizer/code/scripts/test_notion_search.py @@ -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__}")