diff --git a/custom-skills/31-notion-organizer/code/scripts/notion_search.py b/custom-skills/31-notion-organizer/code/scripts/notion_search.py index cf40985..c08ef4a 100644 --- a/custom-skills/31-notion-organizer/code/scripts/notion_search.py +++ b/custom-skills/31-notion-organizer/code/scripts/notion_search.py @@ -11,8 +11,9 @@ from __future__ import annotations import json import re import sys -from typing import Callable, List, Optional +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). @@ -78,3 +79,60 @@ def expand_query( 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 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 b58505c..68322eb 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 @@ -163,6 +163,104 @@ def test_expand_query_falls_back_on_llm_exception(): _assert(variants == ["AI agents"], "LLM exception falls back to original only") +def test_search_unions_and_dedupes_workspace(): + """Two variants returning overlapping pages produce a deduped candidate set.""" + from unittest.mock import MagicMock + import notion_search + + notion = MagicMock() + notion.search.side_effect = [ + {"results": [ + {"id": "page-a", "object": "page", "url": "https://notion.so/a"}, + {"id": "page-b", "object": "page", "url": "https://notion.so/b"}, + ]}, + {"results": [ + {"id": "page-b", "object": "page", "url": "https://notion.so/b"}, + {"id": "page-c", "object": "page", "url": "https://notion.so/c"}, + ]}, + ] + candidates = notion_search.search_candidates( + notion, ["query1", "query2"], databases=None + ) + ids = [c["id"] for c in candidates] + _assert(set(ids) == {"page-a", "page-b", "page-c"}, "union of pages from both variants") + _assert(len(ids) == 3, "duplicate page-b deduped") + + +def test_search_caps_candidates_at_30(): + """If total candidates exceed 30, cap to 30.""" + from unittest.mock import MagicMock + import notion_search + + notion = MagicMock() + notion.search.return_value = { + "results": [ + {"id": f"page-{i:02d}", "object": "page", "url": f"https://notion.so/{i}"} + for i in range(40) + ] + } + candidates = notion_search.search_candidates(notion, ["query"], databases=None) + _assert(len(candidates) == 30, f"capped to 30 (got {len(candidates)})") + + +def test_search_uses_data_source_query_when_databases_specified(): + """--databases flag → per-DB data_sources.query instead of workspace search.""" + from unittest.mock import MagicMock, patch + import notion_search + + notion = MagicMock() + notion.data_sources.query.return_value = { + "results": [{"id": "page-a", "url": "https://notion.so/a"}] + } + + with patch("notion_search.compat.resolve_data_source_id", + side_effect=lambda c, ds: ds): + candidates = notion_search.search_candidates( + notion, ["query"], databases=["db-id-1"] + ) + _assert(notion.data_sources.query.called, "data_sources.query was called") + _assert(not notion.search.called, "workspace search NOT called when databases specified") + _assert(len(candidates) == 1, "candidate from data source returned") + + +def test_search_filters_only_pages(): + """search() may return databases too; we keep only object='page'.""" + from unittest.mock import MagicMock + import notion_search + + notion = MagicMock() + notion.search.return_value = {"results": [ + {"id": "page-a", "object": "page", "url": "https://notion.so/a"}, + {"id": "db-1", "object": "database", "url": "https://notion.so/db1"}, + ]} + candidates = notion_search.search_candidates(notion, ["query"], databases=None) + ids = [c["id"] for c in candidates] + _assert("page-a" in ids, "page kept") + _assert("db-1" not in ids, "database object filtered out") + + +def test_search_passes_property_filter_to_data_sources_query(): + """--filter JSON is passed through to data_sources.query as the `filter` parameter.""" + from unittest.mock import MagicMock, patch + import notion_search + + notion = MagicMock() + notion.data_sources.query.return_value = {"results": []} + + prop_filter = {"property": "Status", "status": {"equals": "Done"}} + + with patch("notion_search.compat.resolve_data_source_id", + side_effect=lambda c, ds: ds): + notion_search.search_candidates( + notion, ["query"], databases=["db-id-1"], prop_filter=prop_filter, + ) + + _assert(notion.data_sources.query.called, "data_sources.query was called") + call_kwargs = notion.data_sources.query.call_args.kwargs + _assert(call_kwargs.get("filter") == prop_filter, + "prop_filter passed through as `filter` keyword argument") + + def run_all(): tests = [ test_call_claude_dispatches_to_sdk_when_available, @@ -177,6 +275,11 @@ def run_all(): test_expand_query_dedupes_and_caps, test_expand_query_falls_back_on_invalid_json, test_expand_query_falls_back_on_llm_exception, + test_search_unions_and_dedupes_workspace, + test_search_caps_candidates_at_30, + test_search_uses_data_source_query_when_databases_specified, + test_search_filters_only_pages, + test_search_passes_property_filter_to_data_sources_query, ] for t in tests: print(f"\n{t.__name__}")