our-claude-skills/custom-skills/31-notion-organizer/code/scripts/notion_search.py

#!/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]
    # Pass cache_dir only if explicitly set; otherwise let _search_cache use its default.
    cache_kwargs = {"cache_dir": cache_dir} if cache_dir else {}

    if use_cache:
        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:
        _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
    if prop_filter and not databases:
        print("Warning: --filter is only applied in per-database mode; ignored without --databases",
              file=sys.stderr)
        prop_filter = 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()