Andrew Yim
5f6438ecf3
Code review polish: - Add stderr warning for the fourth failure path (LLM returned valid JSON but it's not a list-of-strings). Previously this fell back silently, making debugging harder when Haiku produces unexpected shapes. - Drop unused f-string prefixes from two warnings (no interpolation). - Type hint: Callable[..., str] = None → Optional[Callable[..., str]] = None for strict type-checker compatibility. 12/12 tests still pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
#!/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, Optional
|
|
|
|
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
|