Andrew Yim
e67209b905
Code review polish: - Drop the unused `name` parameter from `_flatten_property` (it was reserved for future per-property-name special-casing but never used). - Add a regression test pinning that checkbox=False and number=0 are preserved in the enriched output. The existing empty-value filter is `if value not in (None, [], "")` which keeps falsy-but-meaningful values, but the contract wasn't tested. 22/22 tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
239 lines
8.4 KiB
Python
239 lines
8.4 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 json
|
|
import re
|
|
import sys
|
|
from typing import Callable, Dict, List, Optional
|
|
|
|
import _notion_compat as compat
|
|
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
|