From 9f126fedc7eac1e10699e04dc3a3cd729477ba5e Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Wed, 15 Jul 2026 19:47:43 -0400 Subject: [PATCH] feat(ebook-search): implement async phrase matching for chunks and add ChunkPhraseHit model --- .../protected_phrases/matching.py | 50 ++++++++++++++++++- .../ebook_search/protected_phrases/models.py | 15 ++++++ .../ebook_search/protected_phrases/store.py | 32 ++---------- tests/ebook_search/test_protected_phrases.py | 5 ++ 4 files changed, 74 insertions(+), 28 deletions(-) diff --git a/python/ebook_search/protected_phrases/matching.py b/python/ebook_search/protected_phrases/matching.py index 964d848..c9775c2 100644 --- a/python/ebook_search/protected_phrases/matching.py +++ b/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()} diff --git a/python/ebook_search/protected_phrases/models.py b/python/ebook_search/protected_phrases/models.py index 5213364..164147f 100644 --- a/python/ebook_search/protected_phrases/models.py +++ b/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. diff --git a/python/ebook_search/protected_phrases/store.py b/python/ebook_search/protected_phrases/store.py index 11b8a05..26e8bab 100644 --- a/python/ebook_search/protected_phrases/store.py +++ b/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.sqlite import insert as sqlite_insert +from sqlalchemy.dialects.postgresql import 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, diff --git a/tests/ebook_search/test_protected_phrases.py b/tests/ebook_search/test_protected_phrases.py index f3d89e0..57e18bf 100644 --- a/tests/ebook_search/test_protected_phrases.py +++ b/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")