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