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 46f7db1..dc35602 100644 --- a/custom-skills/31-notion-organizer/code/scripts/notion_search.py +++ b/custom-skills/31-notion-organizer/code/scripts/notion_search.py @@ -8,12 +8,16 @@ CLI: python3 notion_search.py "query" [--databases ID,...] [--filter JSON] \ from __future__ import annotations +import argparse import json +import os import re import sys +from pathlib import Path from typing import Callable, Dict, List, Optional import _notion_compat as compat +import _search_cache from _search_llm import call_claude as default_llm_caller EXPAND_PROMPT = """You are a query expander for a Notion semantic search tool. Generate up to {n} variants of the user's query that capture related concepts, synonyms, and cross-language alternates (especially Korean ↔ English). @@ -330,3 +334,136 @@ def rerank( "snippet": entry.get("why", ""), }) return out + + +def run_search( + notion, + query: str, + *, + databases: Optional[List[str]] = None, + prop_filter: Optional[Dict] = None, + limit: int = 10, + no_rerank: bool = False, + no_expand: bool = False, + use_cache: bool = True, + cache_dir: Optional[Path] = None, + expand_llm: Optional[Callable[..., str]] = None, + rerank_llm: Optional[Callable[..., str]] = None, +) -> List[Dict]: + """Full pipeline: expand → search → enrich → rerank → return. + + expand_llm and rerank_llm are dependency-injected for tests. + Output schema (each result dict): id, url, title, relevance (or None), + snippet (or None), excerpt, properties. + """ + # Stage 1: expand + if no_expand: + queries = [query] + else: + queries = expand_query(query, llm_caller=expand_llm) + + # Stage 2: search + candidates = search_candidates( + notion, queries, databases=databases, prop_filter=prop_filter, + ) + if not candidates: + return [] + + # Stage 3: enrich + enriched = enrich_candidates(notion, candidates) + + # Stage 4: rerank (or skip) + if no_rerank: + return [ + {**c, "relevance": None, "snippet": None} + for c in enriched[:limit] + ] + + candidate_ids = [c["id"] for c in enriched] + if use_cache: + cache_kwargs = {"cache_dir": cache_dir} if cache_dir else {} + cached = _search_cache.cache_get(query, candidate_ids, **cache_kwargs) + if cached is not None: + return cached + + ranked = rerank(query, enriched, llm_caller=rerank_llm, limit=limit) + + if use_cache: + cache_kwargs = {"cache_dir": cache_dir} if cache_dir else {} + _search_cache.cache_put(query, candidate_ids, ranked, **cache_kwargs) + + return ranked + + +def format_terminal(results: List[Dict]) -> str: + """Human-readable terminal output.""" + if not results: + return "No matches.\n" + lines = [] + for i, r in enumerate(results, 1): + rel = f"(rel: {r['relevance']:.2f}) " if r.get("relevance") is not None else "" + lines.append(f"[{i}] {rel}{r['title']}") + props = r.get("properties", {}) + if props: + prop_strs = [] + for k, v in props.items(): + if isinstance(v, list): + v = ", ".join(str(x) for x in v) + prop_strs.append(f"{k}: {v}") + lines.append(f" {' · '.join(prop_strs)}") + if r.get("snippet"): + lines.append(f" Why: {r['snippet']}") + if r.get("url"): + lines.append(f" {r['url']}") + lines.append("") + return "\n".join(lines).rstrip() + "\n" + + +def main(): + parser = argparse.ArgumentParser( + prog="notion-search", + description="Semantic search across Notion workspace via LLM expand + rerank.", + ) + parser.add_argument("query", help="Search query (natural language)") + parser.add_argument("--databases", "-d", default=None, + help="Comma-separated database/data_source IDs (default: workspace-wide)") + parser.add_argument("--filter", "-f", default=None, + help="JSON property filter (per-database mode only)") + parser.add_argument("--limit", "-l", type=int, default=10, + help="Max results after rerank (default: 10)") + parser.add_argument("--no-rerank", action="store_true", + help="Skip Claude rerank stage") + parser.add_argument("--no-expand", action="store_true", + help="Skip query-variant generation") + parser.add_argument("--no-cache", action="store_true", + help="Bypass result cache") + parser.add_argument("--json", action="store_true", + help="Output JSON array instead of terminal table") + args = parser.parse_args() + + # Parse databases and filter + databases = args.databases.split(",") if args.databases else None + prop_filter = json.loads(args.filter) if args.filter else None + + # Build notion client + api_key = os.getenv("NOTION_API_KEY") or os.getenv("NOTION_TOKEN") + if not api_key: + print("Error: NOTION_API_KEY (or NOTION_TOKEN) not set", file=sys.stderr) + sys.exit(1) + notion = compat.make_client(api_key) + + results = run_search( + notion, args.query, + databases=databases, prop_filter=prop_filter, limit=args.limit, + no_rerank=args.no_rerank, no_expand=args.no_expand, + use_cache=not args.no_cache, + ) + + if args.json: + print(json.dumps(results, ensure_ascii=False, indent=2)) + else: + print(format_terminal(results)) + + +if __name__ == "__main__": + main() 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 45464b9..590b11c 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 @@ -457,6 +457,143 @@ def test_rerank_falls_back_on_llm_exception(): _assert(ranked[0]["relevance"] is None, "no score in fallback") +def test_pipeline_search_returns_json_serializable_results(): + """End-to-end pipeline returns results matching the spec's JSON schema.""" + from unittest.mock import MagicMock + import json as json_mod + import notion_search + + notion = MagicMock() + notion.search.return_value = { + "results": [ + { + "id": "page-a", "object": "page", "url": "https://notion.so/a", + "properties": { + "Name": {"type": "title", "title": [{"plain_text": "AI Agents Page"}]}, + "Status": {"type": "status", "status": {"name": "Done"}}, + }, + } + ] + } + notion.blocks.children.list.return_value = { + "results": [ + {"type": "paragraph", + "paragraph": {"rich_text": [{"plain_text": "Notes about AI agents."}]}} + ] + } + + expand_llm = lambda prompt, **kw: '["AI agents"]' + rerank_llm = lambda prompt, **kw: '[{"index": 0, "score": 0.9, "why": "match"}]' + + results = notion_search.run_search( + notion, "AI agents", + expand_llm=expand_llm, rerank_llm=rerank_llm, + limit=10, use_cache=False, + ) + + _assert(len(results) == 1, "one result") + r = results[0] + _assert(r["id"] == "page-a", "id present") + _assert(r["title"] == "AI Agents Page", "title extracted") + _assert(r["relevance"] == 0.9, "relevance score from rerank") + _assert(r["properties"]["Status"] == "Done", "property included") + _assert(r["excerpt"] == "Notes about AI agents.", "excerpt included") + # Schema must be JSON-serializable + json_mod.dumps(results) + + +def test_pipeline_no_rerank_skips_llm(): + """--no-rerank flag → results have null relevance, no rerank LLM call.""" + from unittest.mock import MagicMock + import notion_search + + notion = MagicMock() + notion.search.return_value = { + "results": [ + {"id": "p1", "object": "page", "url": "u1", + "properties": {"Name": {"type": "title", "title": [{"plain_text": "Page 1"}]}}} + ] + } + notion.blocks.children.list.return_value = {"results": []} + + expand_llm = lambda prompt, **kw: '["query"]' + + rerank_calls = [] + def rerank_llm(prompt, **kw): + rerank_calls.append(prompt) + return "[]" + + results = notion_search.run_search( + notion, "query", + expand_llm=expand_llm, rerank_llm=rerank_llm, + no_rerank=True, use_cache=False, + ) + _assert(len(rerank_calls) == 0, "rerank LLM not called when --no-rerank") + _assert(results[0]["relevance"] is None, "no relevance score when no rerank") + + +def test_pipeline_no_expand_uses_single_query(): + """--no-expand → expand_llm not called, Notion gets only original query.""" + from unittest.mock import MagicMock + import notion_search + + notion = MagicMock() + notion.search.return_value = {"results": []} + + expand_calls = [] + def expand_llm(prompt, **kw): + expand_calls.append(prompt) + return '["should not be used"]' + + rerank_llm = lambda prompt, **kw: "[]" + + notion_search.run_search( + notion, "exact term", + expand_llm=expand_llm, rerank_llm=rerank_llm, + no_expand=True, use_cache=False, + ) + _assert(len(expand_calls) == 0, "expand LLM not called when --no-expand") + _assert(notion.search.call_count == 1, "Notion search called exactly once") + + +def test_pipeline_uses_cache_on_second_call(): + """Two identical pipeline calls → second one skips rerank LLM.""" + import tempfile + from pathlib import Path + from unittest.mock import MagicMock + import notion_search + + notion = MagicMock() + notion.search.return_value = { + "results": [ + {"id": "p1", "object": "page", "url": "u1", + "properties": {"Name": {"type": "title", "title": [{"plain_text": "Page 1"}]}}} + ] + } + notion.blocks.children.list.return_value = {"results": []} + + expand_llm = lambda prompt, **kw: '["query"]' + + rerank_calls = [] + def rerank_llm(prompt, **kw): + rerank_calls.append(prompt) + return '[{"index": 0, "score": 0.5, "why": "x"}]' + + with tempfile.TemporaryDirectory() as tmpdir: + cache_dir = Path(tmpdir) + + notion_search.run_search( + notion, "query", expand_llm=expand_llm, rerank_llm=rerank_llm, + use_cache=True, cache_dir=cache_dir, + ) + notion_search.run_search( + notion, "query", expand_llm=expand_llm, rerank_llm=rerank_llm, + use_cache=True, cache_dir=cache_dir, + ) + + _assert(len(rerank_calls) == 1, "second call skipped rerank LLM (cache hit)") + + def run_all(): tests = [ test_call_claude_dispatches_to_sdk_when_available, @@ -485,6 +622,10 @@ def run_all(): test_rerank_respects_limit, test_rerank_falls_back_on_parse_error, test_rerank_falls_back_on_llm_exception, + test_pipeline_search_returns_json_serializable_results, + test_pipeline_no_rerank_skips_llm, + test_pipeline_no_expand_uses_single_query, + test_pipeline_uses_cache_on_second_call, ] for t in tests: print(f"\n{t.__name__}")