From fffbab201d54f5374bd6257445f563c0e7a0d5e1 Mon Sep 17 00:00:00 2001 From: Andrew Yim Date: Tue, 28 Apr 2026 07:47:20 +0900 Subject: [PATCH] feat(notion-search): add query expansion via Claude Haiku MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Generates up to 5 query variants (synonyms + cross-language KR↔EN) so later Notion API search can union over them. Permissive failure modes: LLM error or non-JSON response falls back to [original] with stderr warning. Dedupes and caps variants. 4 tests: variants list, dedup+cap, JSON-parse fallback, exception fallback. 12 passing total. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../code/scripts/notion_search.py | 78 +++++++++++++++++++ .../code/scripts/test_notion_search.py | 45 +++++++++++ 2 files changed, 123 insertions(+) create mode 100644 custom-skills/31-notion-organizer/code/scripts/notion_search.py diff --git a/custom-skills/31-notion-organizer/code/scripts/notion_search.py b/custom-skills/31-notion-organizer/code/scripts/notion_search.py new file mode 100644 index 0000000..57b7240 --- /dev/null +++ b/custom-skills/31-notion-organizer/code/scripts/notion_search.py @@ -0,0 +1,78 @@ +#!/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, List + +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: 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(f"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(f"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): + 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 diff --git a/custom-skills/31-notion-organizer/code/scripts/test_notion_search.py b/custom-skills/31-notion-organizer/code/scripts/test_notion_search.py index 1f0f2b5..b58505c 100644 --- a/custom-skills/31-notion-organizer/code/scripts/test_notion_search.py +++ b/custom-skills/31-notion-organizer/code/scripts/test_notion_search.py @@ -122,6 +122,47 @@ def test_cache_handles_corrupted_file(): _assert(result is None, "corrupted cache file returns None") +def test_expand_query_returns_variants_with_original_first(): + """LLM mock returns variants; expansion includes original verbatim as first item.""" + import notion_search + + fake_llm = lambda prompt, **kw: '["AI agents", "multi-agent systems", "autonomous LLMs"]' + variants = notion_search.expand_query("AI agents", llm_caller=fake_llm) + _assert(variants[0] == "AI agents", "original query is first variant") + _assert(len(variants) >= 2, "at least 2 variants returned") + _assert("multi-agent systems" in variants, "LLM-suggested variant included") + + +def test_expand_query_dedupes_and_caps(): + """Duplicates removed; total ≤ max_variants.""" + import notion_search + + fake_llm = lambda prompt, **kw: '["AI", "AI", "agents", "agents", "systems", "tools", "models"]' + variants = notion_search.expand_query("AI", llm_caller=fake_llm, max_variants=3) + _assert(len(variants) <= 3, "capped at max_variants=3") + _assert(len(variants) == len(set(variants)), "no duplicates in result") + + +def test_expand_query_falls_back_on_invalid_json(): + """LLM returns prose instead of JSON → fall back to [original].""" + import notion_search + + fake_llm = lambda prompt, **kw: "I'm sorry, I can't help with that." + variants = notion_search.expand_query("AI agents", llm_caller=fake_llm) + _assert(variants == ["AI agents"], "non-JSON response falls back to original only") + + +def test_expand_query_falls_back_on_llm_exception(): + """LLM raises → fall back to [original].""" + import notion_search + + def boom(prompt, **kw): + raise RuntimeError("API error") + + variants = notion_search.expand_query("AI agents", llm_caller=boom) + _assert(variants == ["AI agents"], "LLM exception falls back to original only") + + def run_all(): tests = [ test_call_claude_dispatches_to_sdk_when_available, @@ -132,6 +173,10 @@ def run_all(): test_cache_miss_on_different_candidates, test_cache_miss_after_ttl_expires, test_cache_handles_corrupted_file, + test_expand_query_returns_variants_with_original_first, + test_expand_query_dedupes_and_caps, + test_expand_query_falls_back_on_invalid_json, + test_expand_query_falls_back_on_llm_exception, ] for t in tests: print(f"\n{t.__name__}")