feat(ebook-search): implement async phrase matching for chunks and add ChunkPhraseHit model

python/ebook_search/protected_phrases/matching.py

@@ -6,10 +6,11 @@ import logging
 from collections import defaultdict
 from typing import TYPE_CHECKING
 
-from sqlalchemy import and_, delete, or_, select, union
+from sqlalchemy import and_, delete, func, or_, select, union
 
 from python.ebook_search.protected_phrases.config import get_ignored_phrases
 from python.ebook_search.protected_phrases.models import (
+    ChunkPhraseHit,
     PhraseLookup,
     PhraseMatch,
 )
@@ -355,3 +356,50 @@ def generate_query_ngrams(
             if phrase_norm in get_ignored_phrases():
                 continue
             yield phrase_norm, start, end
+
+
+async def phrase_hits_for_chunks(
+    session: AsyncSession,
+    *,
+    chunk_ids: Sequence[int],
+    phrase_ids: Sequence[int],
+) -> dict[int, tuple[ChunkPhraseHit, ...]]:
+    """Return matched protected phrases with mention counts by chunk id.
+
+    Args:
+        session (AsyncSession): Active database session.
+        chunk_ids (Sequence[int]): Chunk ids to look up mentions for.
+        phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to.
+
+    Returns:
+        dict[int, tuple[ChunkPhraseHit, ...]]: Phrase hits per chunk id, ordered by mention count.
+    """
+    if not chunk_ids or not phrase_ids:
+        return {}
+
+    mention_count = func.count(EbookChunkPhraseMention.phrase_id).label("mention_count")
+    statement = (
+        select(
+            EbookChunkPhraseMention.chunk_id,
+            EbookProtectedPhrase.id.label("phrase_id"),
+            EbookProtectedPhrase.phrase_text,
+            mention_count,
+        )
+        .join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookChunkPhraseMention.phrase_id)
+        .where(
+            EbookChunkPhraseMention.chunk_id.in_(chunk_ids),
+            EbookChunkPhraseMention.phrase_id.in_(phrase_ids),
+        )
+        .group_by(EbookChunkPhraseMention.chunk_id, EbookProtectedPhrase.id, EbookProtectedPhrase.phrase_text)
+        .order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text)
+    )
+    hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list)
+    for row in await session.execute(statement):
+        hits[row.chunk_id].append(
+            ChunkPhraseHit(
+                phrase_id=row.phrase_id,
+                phrase_text=row.phrase_text,
+                mention_count=row.mention_count,
+            )
+        )
+    return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()}

python/ebook_search/protected_phrases/models.py

@@ -131,6 +131,21 @@ class PhraseMatch:
     series_id: int | None = None
 
 
+@dataclass(frozen=True, slots=True)
+class ChunkPhraseHit:
+    """One protected phrase with its mention count inside one retrieved chunk.
+
+    Attributes:
+        phrase_id (int): Protected phrase id.
+        phrase_text (str): Display text for the phrase.
+        mention_count (int): Indexed mentions inside the chunk.
+    """
+
+    phrase_id: int
+    phrase_text: str
+    mention_count: int
+
+
 @dataclass(frozen=True, slots=True)
 class PhraseCandidateGenerationResult:
     """Summary of candidate phrase extraction for indexed books.

python/ebook_search/protected_phrases/store.py

@@ -7,8 +7,7 @@ import re
 from typing import TYPE_CHECKING
 
 from sqlalchemy import delete, func, or_, select
-from sqlalchemy.dialects.postgresql import insert as pg_insert
+from sqlalchemy.dialects.postgresql import insert
-from sqlalchemy.dialects.sqlite import insert as sqlite_insert
 
 from python.ebook_search.protected_phrases.extraction import minimum_candidate_raw_count
 from python.ebook_search.protected_phrases.models import (
@@ -29,35 +28,14 @@ from python.orm.richie import (
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    from sqlalchemy.dialects.postgresql.dml import Insert as PostgresInsert
-    from sqlalchemy.dialects.sqlite.dml import Insert as SqliteInsert
     from sqlalchemy.ext.asyncio import AsyncSession
 
     from python.ebook_search.config import EbookSearchConfig
     from python.ebook_search.protected_phrases.models import LLMJudgment
-    from python.orm.richie.base import TableBase
 
 logger = logging.getLogger(__name__)
 
 
-def dialect_insert(session: AsyncSession, table: type[TableBase]) -> PostgresInsert | SqliteInsert:
-    """Return a dialect-specific INSERT construct that supports ``ON CONFLICT DO UPDATE``.
-
-    Production runs on PostgreSQL while tests run on SQLite; both support upserts with
-    compatible SQLAlchemy constructs, so the correct one is chosen from the bound dialect.
-
-    Args:
-        session (AsyncSession): Active database session whose bind selects the dialect.
-        table (type[TableBase]): Mapped table to insert into.
-
-    Returns:
-        PostgresInsert | SqliteInsert: A dialect insert exposing ``on_conflict_do_update``.
-    """
-    if session.get_bind().dialect.name == "sqlite":
-        return sqlite_insert(table)
-    return pg_insert(table)
-
-
 async def load_book_text(session: AsyncSession, book_id: int) -> str:
     """Load a book's indexed chunk text as one string for phrase extraction.
 
@@ -377,7 +355,7 @@ async def save_candidate_to_db(
     skip_update = {"book_id", "phrase_norm"}
     if judgment is None:
         skip_update.add("llm_judged")
-    insert_statement = dialect_insert(session, EbookCandidatePhrase).values(**values)
+    insert_statement = insert(EbookCandidatePhrase).values(**values)
     statement = insert_statement.on_conflict_do_update(
         index_elements=["book_id", "phrase_norm"],
         set_={column: insert_statement.excluded[column] for column in values if column not in skip_update},
@@ -423,7 +401,7 @@ async def bulk_upsert_unjudged_candidates(
     skip_update = {"book_id", "phrase_norm", "llm_judged"}
     for chunk_start in range(0, len(values), BULK_CANDIDATE_UPSERT_CHUNK):
         chunk = values[chunk_start : chunk_start + BULK_CANDIDATE_UPSERT_CHUNK]
-        insert_statement = dialect_insert(session, EbookCandidatePhrase).values(chunk)
+        insert_statement = insert(EbookCandidatePhrase).values(chunk)
         statement = insert_statement.on_conflict_do_update(
             index_elements=["book_id", "phrase_norm"],
             set_={column: insert_statement.excluded[column] for column in chunk[0] if column not in skip_update},
@@ -509,7 +487,7 @@ async def upsert_protected_phrase(
         "suppress_children": judgment.suppress_children,
         "source_candidate_id": source_candidate.id,
     }
-    insert_statement = dialect_insert(session, EbookProtectedPhrase).values(**values)
+    insert_statement = insert(EbookProtectedPhrase).values(**values)
     statement = insert_statement.on_conflict_do_update(
         index_elements=["book_id", "phrase_norm"],
         set_={
@@ -542,7 +520,7 @@ async def upsert_phrase_alias(
     if not alias_norm or alias_norm == phrase.phrase_norm:
         return None
 
-    insert_statement = dialect_insert(session, EbookPhraseAlias).values(
+    insert_statement = insert(EbookPhraseAlias).values(
         phrase_id=phrase.id,
         alias_norm=alias_norm,
         alias_text=alias_text,

tests/ebook_search/test_protected_phrases.py

@@ -33,9 +33,11 @@ from python.ebook_search.protected_phrases.matching import (
     detect_protected_phrases_for_query,
     index_chunk_phrase_mentions,
     load_phrase_lookup,
+    phrase_hits_for_chunks,
     resolve_overlaps,
 )
 from python.ebook_search.protected_phrases.models import (
+    ChunkPhraseHit,
     LLMJudgment,
     PhraseCandidate,
     PhraseMatch,
@@ -353,6 +355,9 @@ async def test_index_chunk_phrase_mentions_uses_normalized_window_lookup(
     assert mention.chunk_id == chunk.id
     assert mention.phrase_id == phrase.id
     assert chunk.text[mention.start_char : mention.end_char] == "lock-in"
+    assert await phrase_hits_for_chunks(session, chunk_ids=[chunk.id], phrase_ids=[phrase.id]) == {
+        chunk.id: (ChunkPhraseHit(phrase_id=phrase.id, phrase_text="lock in", mention_count=1),)
+    }
 
 
 @pytest.mark.usefixtures("worker_pool")