feat(notion-search): add CLI entrypoint with argparse + output formatting

Wires together the four stages (expand → search → enrich → rerank) into
run_search(). CLI flags: --databases, --filter, --limit, --no-rerank,
--no-expand, --no-cache, --json. Terminal output renders as a numbered
table with title, relevance, properties, snippet, URL.

Cache lookup happens BEFORE rerank, with cache_put after success.
NOTION_API_KEY (or NOTION_TOKEN) env var required.

4 end-to-end pipeline tests (mocked Notion + LLM): JSON-serializable
output, --no-rerank skip, --no-expand skip, cache hit on repeat.

Total: 30 tests passing.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-04-28 08:07:00 +09:00
parent 20029ecc9c
commit 72d4b36943
2 changed files with 278 additions and 0 deletions

View File

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

View File

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