Andrew Yim
72d4b36943
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>
470 lines
16 KiB
Python
470 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Notion semantic search — query expansion + LLM rerank over Notion API search.
|
|
|
|
CLI: python3 notion_search.py "query" [--databases ID,...] [--filter JSON] \
|
|
[--limit N] [--no-rerank] [--no-expand] \
|
|
[--no-cache] [--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).
|
|
|
|
Rules:
|
|
- Always include the original query verbatim as the first variant.
|
|
- Variants should help find pages that the original query might miss due to keyword-only search.
|
|
- For Korean queries, include English synonyms; for English queries, include Korean alternates if the topic has common Korean usage.
|
|
- Keep variants concise (under 8 words each).
|
|
|
|
Query: {query}
|
|
|
|
Return ONLY a JSON array of strings, no prose. Example:
|
|
["original query", "synonym variant", "cross-language variant"]"""
|
|
|
|
|
|
def expand_query(
|
|
query: str,
|
|
*,
|
|
llm_caller: Optional[Callable[..., str]] = None,
|
|
max_variants: int = 5,
|
|
) -> List[str]:
|
|
"""Expand a query into related variants. Returns [query] on any failure."""
|
|
if llm_caller is None:
|
|
llm_caller = default_llm_caller
|
|
|
|
prompt = EXPAND_PROMPT.format(n=max_variants, query=query)
|
|
|
|
try:
|
|
response = llm_caller(prompt)
|
|
except Exception as exc:
|
|
print(f"Warning: query expansion failed ({exc}); using original query only",
|
|
file=sys.stderr)
|
|
return [query]
|
|
|
|
# Be permissive: extract the first JSON array from the response.
|
|
match = re.search(r'\[.*\]', response, re.DOTALL)
|
|
if not match:
|
|
print("Warning: query expansion returned no JSON array; using original query only",
|
|
file=sys.stderr)
|
|
return [query]
|
|
|
|
try:
|
|
variants = json.loads(match.group(0))
|
|
except json.JSONDecodeError:
|
|
print("Warning: query expansion returned invalid JSON; using original query only",
|
|
file=sys.stderr)
|
|
return [query]
|
|
|
|
if not isinstance(variants, list) or not all(isinstance(v, str) for v in variants):
|
|
print("Warning: query expansion returned non-string-list shape; using original query only",
|
|
file=sys.stderr)
|
|
return [query]
|
|
|
|
# Always include the original query first; dedupe; cap at max_variants.
|
|
seen = set()
|
|
result = []
|
|
for v in [query] + variants:
|
|
v = v.strip()
|
|
if v and v not in seen:
|
|
seen.add(v)
|
|
result.append(v)
|
|
if len(result) >= max_variants:
|
|
break
|
|
return result
|
|
|
|
|
|
MAX_CANDIDATES = 30
|
|
|
|
|
|
def search_candidates(
|
|
notion,
|
|
queries: List[str],
|
|
*,
|
|
databases: Optional[List[str]] = None,
|
|
prop_filter: Optional[Dict] = None,
|
|
) -> List[Dict]:
|
|
"""Run each query against Notion (workspace or per-DB), dedupe, cap at MAX_CANDIDATES.
|
|
|
|
Returns a list of page result objects (whatever Notion returned), order preserves
|
|
first-seen position across all variants — used as fallback ordering when rerank is off.
|
|
"""
|
|
seen_ids = set()
|
|
candidates: List[Dict] = []
|
|
|
|
for query in queries:
|
|
if databases:
|
|
# Per-database query (data_sources.query)
|
|
for db_id in databases:
|
|
try:
|
|
data_source_id = compat.resolve_data_source_id(notion, db_id)
|
|
except Exception as exc:
|
|
print(f"Warning: could not resolve database {db_id}: {exc}",
|
|
file=sys.stderr)
|
|
continue
|
|
|
|
params = {"data_source_id": data_source_id, "page_size": 100}
|
|
if prop_filter:
|
|
params["filter"] = prop_filter
|
|
|
|
try:
|
|
response = notion.data_sources.query(**params)
|
|
except Exception as exc:
|
|
print(f"Warning: query failed for database {db_id}: {exc}",
|
|
file=sys.stderr)
|
|
continue
|
|
for page in response.get("results") or []:
|
|
page_id = page.get("id")
|
|
if page_id and page_id not in seen_ids:
|
|
seen_ids.add(page_id)
|
|
candidates.append(page)
|
|
if len(candidates) >= MAX_CANDIDATES:
|
|
return candidates
|
|
else:
|
|
# Workspace-wide search
|
|
try:
|
|
response = notion.search(query=query, page_size=100)
|
|
except Exception as exc:
|
|
print(f"Warning: workspace search failed for variant {query!r}: {exc}",
|
|
file=sys.stderr)
|
|
continue
|
|
for item in response.get("results") or []:
|
|
if item.get("object") != "page":
|
|
continue
|
|
page_id = item.get("id")
|
|
if page_id and page_id not in seen_ids:
|
|
seen_ids.add(page_id)
|
|
candidates.append(item)
|
|
if len(candidates) >= MAX_CANDIDATES:
|
|
return candidates
|
|
|
|
return candidates
|
|
|
|
|
|
EXCERPT_MAX_CHARS = 200
|
|
TEXT_BEARING_BLOCK_TYPES = {"paragraph", "heading_1", "heading_2", "heading_3",
|
|
"quote", "callout", "bulleted_list_item",
|
|
"numbered_list_item", "to_do", "toggle"}
|
|
|
|
|
|
def _flatten_property(prop: Dict):
|
|
"""Flatten a Notion property to a Python value suitable for display/rerank."""
|
|
ptype = prop.get("type")
|
|
if ptype == "title":
|
|
return "".join(t.get("plain_text", "") for t in prop.get("title", []))
|
|
if ptype == "rich_text":
|
|
return "".join(t.get("plain_text", "") for t in prop.get("rich_text", []))
|
|
if ptype == "select":
|
|
sel = prop.get("select")
|
|
return sel.get("name") if sel else None
|
|
if ptype == "status":
|
|
st = prop.get("status")
|
|
return st.get("name") if st else None
|
|
if ptype == "multi_select":
|
|
return [o.get("name") for o in prop.get("multi_select", [])]
|
|
if ptype == "date":
|
|
return prop.get("date")
|
|
if ptype == "checkbox":
|
|
return prop.get("checkbox")
|
|
if ptype == "number":
|
|
return prop.get("number")
|
|
if ptype == "url":
|
|
return prop.get("url")
|
|
if ptype == "email":
|
|
return prop.get("email")
|
|
if ptype == "phone_number":
|
|
return prop.get("phone_number")
|
|
return None
|
|
|
|
|
|
def _extract_title(properties: Dict) -> str:
|
|
for prop in properties.values():
|
|
if prop.get("type") == "title":
|
|
return "".join(t.get("plain_text", "") for t in prop.get("title", []))
|
|
return ""
|
|
|
|
|
|
def _block_text(block: Dict) -> str:
|
|
"""Concatenate plain_text from a block's rich_text array, if any."""
|
|
btype = block.get("type")
|
|
if btype not in TEXT_BEARING_BLOCK_TYPES:
|
|
return ""
|
|
body = block.get(btype, {})
|
|
rich_text = body.get("rich_text", [])
|
|
return "".join(t.get("plain_text", "") for t in rich_text)
|
|
|
|
|
|
def _fetch_excerpt(notion, page_id: str) -> str:
|
|
"""Fetch first text-bearing block; return its plain text capped at EXCERPT_MAX_CHARS."""
|
|
try:
|
|
response = notion.blocks.children.list(block_id=page_id, page_size=5)
|
|
except Exception:
|
|
return ""
|
|
for block in response.get("results") or []:
|
|
text = _block_text(block).strip()
|
|
if text:
|
|
return text[:EXCERPT_MAX_CHARS]
|
|
return ""
|
|
|
|
|
|
def enrich_candidates(notion, candidates: List[Dict]) -> List[Dict]:
|
|
"""Add title, flattened properties, and 200-char excerpt to each candidate."""
|
|
enriched = []
|
|
for c in candidates:
|
|
properties = c.get("properties", {})
|
|
title = _extract_title(properties)
|
|
flat_props = {}
|
|
for name, prop in properties.items():
|
|
if prop.get("type") == "title":
|
|
continue
|
|
value = _flatten_property(prop)
|
|
if value not in (None, [], ""):
|
|
flat_props[name] = value
|
|
excerpt = _fetch_excerpt(notion, c["id"])
|
|
enriched.append({
|
|
"id": c["id"],
|
|
"url": c.get("url", ""),
|
|
"title": title,
|
|
"properties": flat_props,
|
|
"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
|
|
|
|
|
|
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()
|