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 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

                response = notion.data_sources.query(**params)
                for page in response.get("results", []):
                    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
            response = notion.search(query=query, page_size=100)
            for item in response.get("results", []):
                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