#!/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