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

                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