refactor(ebook-search): simplify search and phrase matching

python/ebook_search/search.py

@@ -19,6 +19,7 @@ from python.ebook_search.bm25_corpus import (
     load_bm25_corpus,
     score_bm25_corpus,
 )
+from python.ebook_search.chunk_records import CHUNK_RECORD_COLUMNS
 from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table
 from python.ebook_search.protected_phrases.matching import (
     detect_protected_phrases_for_query,
@@ -40,7 +41,7 @@ if TYPE_CHECKING:
     from sqlalchemy.ext.asyncio import AsyncEngine
 
     from python.ebook_search.config import EbookSearchConfig
-    from python.ebook_search.protected_phrases.models import HydratedPhraseMatch
+    from python.ebook_search.protected_phrases.models import PhraseMatch
 
 logger = logging.getLogger(__name__)
 
@@ -74,7 +75,7 @@ class SearchResponse:
     results: list[SearchResult]
     rank_label: str
     timings: tuple[RuntimeStep, ...] = ()
-    phrase_matches: tuple[HydratedPhraseMatch, ...] = ()
+    phrase_matches: tuple[PhraseMatch, ...] = ()
 
     @property
     def total_runtime_ms(self) -> float:
@@ -88,7 +89,7 @@ class RetrievalResponse:
 
     vector_results: list[SearchResult]
     lexical_results: list[SearchResult]
-    phrase_matches: list[HydratedPhraseMatch]
+    phrase_matches: list[PhraseMatch]
     timings: tuple[RuntimeStep, ...]
 
 
@@ -151,18 +152,17 @@ async def search_ebooks(
     return response
 
 
-def skip_phrase_matches() -> list[HydratedPhraseMatch]:
-    """Return no protected phrase matches when phrase matching is disabled."""
-    logger.info("ebook_protected_phrase_detection_skipped")
-    return []
-
-
 async def query_phrase_matches(
     engine: AsyncEngine,
     query: str,
     config: EbookSearchConfig,
-) -> list[HydratedPhraseMatch]:
+    *,
+    phrase_matching: bool,
+) -> list[PhraseMatch]:
     """Detect protected phrases in a query without making search fail when phrase tables are unavailable."""
+    if not phrase_matching:
+        logger.info("ebook_protected_phrase_detection_skipped")
+        return []
     try:
         async with AsyncSession(engine) as session:
             return await detect_protected_phrases_for_query(session, query, config)
@@ -180,7 +180,7 @@ def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchRes
 async def apply_phrase_mention_boosts(
     engine: AsyncEngine,
     candidates: list[SearchResult],
-    phrase_matches: Sequence[HydratedPhraseMatch],
+    phrase_matches: Sequence[PhraseMatch],
     phrase_hit_boost: float,
 ) -> list[SearchResult]:
     """Boost retrieved chunks that have indexed mentions for detected protected phrases."""
@@ -242,23 +242,21 @@ async def parallel_retrieval(
 
     BM25 scoring is pure CPU work over the cached corpus, so it runs in a worker thread
     instead of on the event loop. Protected phrase detection only depends on the query, so
-    it joins the gather as a third task when phrase matching is enabled.
+    it joins the gather as a third task and returns immediately when phrase matching is disabled.
     """
-    phrase_task = (
-        asyncio.create_task(
-            async_timed_result("Protected phrase detection", query_phrase_matches(engine, query, config))
-        )
-        if phrase_matching
-        else None
-    )
-    (vector_results, vector_timing), (lexical_results, lexical_timing) = await asyncio.gather(
+    phrase_timing_name = "Protected phrase detection" if phrase_matching else "Protected phrase detection skipped"
+    (
+        (vector_results, vector_timing),
+        (lexical_results, lexical_timing),
+        (phrase_matches, phrase_timing),
+    ) = await asyncio.gather(
         async_timed_result("Embedding + vector search", vector_candidates(engine, client, query, config)),
         async_timed_result("BM25 search", asyncio.to_thread(bm25_candidates, query, config)),
+        async_timed_result(
+            phrase_timing_name,
+            query_phrase_matches(engine, query, config, phrase_matching=phrase_matching),
+        ),
     )
-    if phrase_task is not None:
-        phrase_matches, phrase_timing = await phrase_task
-    else:
-        phrase_matches, phrase_timing = timed_result("Protected phrase detection skipped", skip_phrase_matches)
 
     logger.info(
         f"ebook_parallel_retrieval_complete vector_candidates={len(vector_results)} "
@@ -334,13 +332,7 @@ async def vector_candidates(
         score = (literal(1.0) - distance).label("score")
         statement = (
             select(
-                EbookChunk.id.label("chunk_id"),
-                EbookChunk.text.label("text"),
-                EbookSource.id.label("source_id"),
-                EbookSource.title.label("source_title"),
-                EbookSource.author.label("source_author"),
-                EbookChapter.title.label("chapter_title"),
-                EbookChunk.page_label.label("page_label"),
+                *CHUNK_RECORD_COLUMNS,
                 score,
             )
             .select_from(embedding_table)