diff --git a/python/ebook_search/api/routes/admin.py b/python/ebook_search/api/routes/admin.py index 84769ba..68ff64e 100644 --- a/python/ebook_search/api/routes/admin.py +++ b/python/ebook_search/api/routes/admin.py @@ -61,50 +61,20 @@ async def scan_library(request: Request, config: AppConfig, session: AsyncDbSess @router.post("/phrases/generate-all", response_class=HTMLResponse) -async def generate_all_phrases(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse: +async def generate_all_phrases(request: Request, config: AppConfig, engine: AppEngine) -> HTMLResponse: """Regenerate candidate phrases for every indexed book without LLM judging.""" - return await run_phrase_generation(request, config, session) - - -async def run_phrase_generation( - request: Request, - config: AppConfig, - session: AsyncDbSession, - *, - only_missing: bool, -) -> HTMLResponse: - """Run candidate phrase generation and render the outcome as an admin status partial. - - Args: - request (Request): Current request, for template rendering. - config (AppConfig): Runtime phrase-tuning settings. - session (AsyncDbSession): Active database session. - only_missing (bool): Only generate for books without candidates instead of every book. - - Returns: - HTMLResponse: Status partial describing the generation outcome. - """ try: - result = await generate_candidate_phrases_for_books(session, config, only_missing=only_missing) - await session.commit() + result = await generate_candidate_phrases_for_books(engine, config) except Exception as error: - await session.rollback() - logger.exception("ebook_admin_generate_phrases_failed only_missing=%s", only_missing) + logger.exception("ebook_admin_generate_phrases_failed") return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) logger.info( - "ebook_admin_generate_phrases_complete only_missing=%s books_seen=%s books_built=%s candidates=%s", - only_missing, + "ebook_admin_generate_phrases_complete books_seen=%s books_built=%s candidates=%s", result.books_seen, result.books_built, result.candidate_phrases, ) - if only_missing and result.books_seen == 0: - return templates.TemplateResponse( - request, - "partials/admin_status.html", - {"message": "All books already have candidate phrases"}, - ) return templates.TemplateResponse( request, "partials/admin_status.html", diff --git a/python/ebook_search/api/routes/page.py b/python/ebook_search/api/routes/page.py index 47858ee..95fa7e7 100644 --- a/python/ebook_search/api/routes/page.py +++ b/python/ebook_search/api/routes/page.py @@ -168,7 +168,10 @@ async def recalculate_book_phrases(source_id: int, config: AppConfig, session: A if source is None: raise HTTPException(status_code=404, detail="Book not found") - result = await recalculate_candidate_phrases_for_book(session, source, config, use_process_pool=True) + try: + result = await recalculate_candidate_phrases_for_book(session, source, config) + except ValueError as error: + raise HTTPException(status_code=409, detail=str(error)) from error logger.info( "ebook_book_phrase_recalculation_complete source_id=%s candidates=%s deleted_candidates=%s " "deleted_protected=%s deleted_aliases=%s deleted_mentions=%s", diff --git a/python/ebook_search/protected_phrases/generate_ngrams.py b/python/ebook_search/protected_phrases/generate_ngrams.py index 165d9b5..8e2ed04 100644 --- a/python/ebook_search/protected_phrases/generate_ngrams.py +++ b/python/ebook_search/protected_phrases/generate_ngrams.py @@ -4,11 +4,11 @@ from __future__ import annotations import asyncio import logging -from collections import deque from time import perf_counter from typing import TYPE_CHECKING from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book from python.ebook_search.protected_phrases.models import ( @@ -16,65 +16,93 @@ from python.ebook_search.protected_phrases.models import ( PhraseCandidateGenerationResult, PhraseRecalculationResult, ) -from python.ebook_search.protected_phrases.pool import extract_phrase_candidates_in_pool, get_extraction_pool +from python.ebook_search.protected_phrases.pool import get_extraction_pool from python.ebook_search.protected_phrases.store import ( bulk_upsert_unjudged_candidates, delete_phrase_data_for_book, load_book_chapter_texts, - metadata_for_source, + metadata_for_source_id, new_candidate_row, prune_unstorable_unjudged_candidate_phrases, ) -from python.orm.richie import EbookCandidatePhrase, EbookSource +from python.orm.common import get_async_postgres_engine +from python.orm.richie import EbookSource if TYPE_CHECKING: - from collections.abc import Mapping, Sequence - from concurrent.futures import Future - - from sqlalchemy.ext.asyncio import AsyncSession + from sqlalchemy.ext.asyncio import AsyncEngine from python.ebook_search.config import EbookSearchConfig - from python.ebook_search.protected_phrases.extraction import SpacyLanguage from python.ebook_search.protected_phrases.models import PhraseCandidate logger = logging.getLogger(__name__) +class BookHasNoChaptersError(ValueError): + """Raised when a book has no indexed chapter text to generate phrases from.""" + + async def generate_candidate_phrases_for_books( - session: AsyncSession, + engine: AsyncEngine, config: EbookSearchConfig, - *, - only_missing: bool = False, ) -> PhraseCandidateGenerationResult: """Create or refresh candidate phrases for indexed books without calling the LLM judge. - Extraction always runs concurrently in the shared process pool so a full backfill uses - multiple cores. + Every book is submitted to the shared process pool up front and runs in parallel across the + pool's workers; the call blocks until all books have finished. Each worker opens its own + database engine from environment variables, loads the book's chapters, and commits the + book's candidates independently. Args: - session (Session): Active database session. + engine (AsyncEngine): Engine used to read the book list in this process. config (EbookSearchConfig): Runtime phrase-tuning settings. - only_missing (bool): When True, only generate for books that have no candidate phrases - yet instead of refreshing every book. Returns: PhraseCandidateGenerationResult: Per-corpus counts of books seen, built, and candidates stored. + + Results are collected in book order while the pool keeps working. A book failure (including + a book with no indexed chapters) is logged and counted as not built; the remaining books + are unaffected. """ - source_query = select(EbookSource).order_by(EbookSource.id) - if only_missing: - has_candidates = select(EbookCandidatePhrase.id).where(EbookCandidatePhrase.book_id == EbookSource.id) - source_query = source_query.where(~has_candidates.exists()) - sources = (await session.scalars(source_query)).all() - books_seen = len(sources) + async with AsyncSession(engine, expire_on_commit=False) as session: + source_query = select(EbookSource.id).order_by(EbookSource.id) + source_ids = (await session.scalars(source_query)).all() + books_seen = len(source_ids) logger.info( - "ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s max_candidates_per_book=%s", + "ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s " + "max_candidates_per_book=%s", books_seen, config.phrase_min_tokens, config.phrase_max_tokens, config.protected_phrase_max_candidates_per_book, ) - outcomes = await generate_candidates_for_sources_pooled(session, sources, config) + pool = get_extraction_pool(config.protected_phrase_extraction_workers) + wrapped_futures = [ + ( + source_id, + asyncio.wrap_future(pool.submit(generate_candidate_phrases_for_book_in_worker, source_id, None, config)), + ) + for source_id in source_ids + ] + outcomes: list[BookCandidateResult] = [] + for source_id, wrapped_future in wrapped_futures: + await asyncio.wait([wrapped_future]) + exception = wrapped_future.exception() + if exception is not None: + logger.error( + "ebook_candidate_phrase_generation_book_failed source_id=%s", + source_id, + exc_info=exception, + ) + outcomes.append(BookCandidateResult()) + continue + saved_count = wrapped_future.result() + logger.info( + "ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s", + source_id, + saved_count, + ) + outcomes.append(BookCandidateResult(candidates=saved_count, built=True)) result = PhraseCandidateGenerationResult( books_seen=books_seen, @@ -90,113 +118,27 @@ async def generate_candidate_phrases_for_books( return result -async def generate_candidates_for_sources_pooled( - session: AsyncSession, - sources: Sequence[EbookSource], - config: EbookSearchConfig, -) -> list[BookCandidateResult]: - """Generate candidate phrases for many books, extracting them concurrently in worker processes. - - Chapter loading and row persistence stay on the caller's session (serial), while the CPU-bound - extraction runs in the shared process pool. A bounded window of in-flight books overlaps - extraction across cores without loading every book's candidates into memory at once. - - Args: - session (Session): Active database session. - sources (Sequence[EbookSource]): Indexed books to generate candidates for. - config (EbookSearchConfig): Runtime phrase-tuning settings. - - Returns: - list[BookCandidateResult]: One result per book. - """ - pool = get_extraction_pool(config.protected_phrase_extraction_workers) - max_in_flight = max(1, config.protected_phrase_extraction_workers) * 2 - pending: deque[tuple[EbookSource, Future[list[PhraseCandidate]]]] = deque() - outcomes: list[BookCandidateResult] = [] - - async def drain_one() -> None: - source, future = pending.popleft() - extracted = await asyncio.wrap_future(future) - outcomes.append(await store_source_candidates(session, source, extracted, config)) - - try: - for source in sources: - chapters = await load_book_chapter_texts(session, source.id) - if not chapters: - logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id) - outcomes.append(BookCandidateResult()) - continue - future = pool.submit( - extract_phrase_candidates_for_book, - "\n\n".join(chapters), - chapters, - config, - metadata=metadata_for_source(source), - ) - pending.append((source, future)) - if len(pending) >= max_in_flight: - await drain_one() - while pending: - await drain_one() - except Exception: - for _, future in pending: - future.cancel() - await session.rollback() - logger.exception("ebook_candidate_phrase_generation_pooled_failed") - raise - return outcomes - - -async def store_source_candidates( - session: AsyncSession, - source: EbookSource, - limited_candidates: list[PhraseCandidate], - config: EbookSearchConfig, -) -> BookCandidateResult: - """Persist and commit one book's already-extracted candidates. - - Args: - session (AsyncSession): Active database session. - source (EbookSource): Book the candidates belong to. - limited_candidates (list[PhraseCandidate]): Scored candidates to persist. - config (EbookSearchConfig): Runtime phrase-tuning settings. - - Returns: - BookCandidateResult: Candidate count and that the book was committed. - """ - book_started_at = perf_counter() - saved_count = await store_candidate_phrases_for_book(session, source.id, None, limited_candidates, config) - await session.commit() - logger.info( - "ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s duration_ms=%.1f", - source.id, - saved_count, - (perf_counter() - book_started_at) * 1000, - ) - return BookCandidateResult(candidates=saved_count, built=True) - - async def recalculate_candidate_phrases_for_book( session: AsyncSession, source: EbookSource, config: EbookSearchConfig, - *, - nlp: SpacyLanguage | None = None, - use_process_pool: bool = False, ) -> PhraseRecalculationResult: """Remove all book phrase data, regenerate candidates, and commit the completed book. Args: - session (Session): Active database session. + session (AsyncSession): Active database session; deletion and regeneration commit on it. source (EbookSource): Indexed book to recalculate. config (EbookSearchConfig): Runtime phrase-tuning settings. - nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. - use_process_pool (bool): Run the CPU-bound extraction in a worker process so concurrent - recalculations do not serialize behind the GIL. Defaults to in-process for callers - (tests, backfills) that do not need it. Returns: PhraseRecalculationResult: Deleted-row counts and the number of candidates regenerated. + + Raises: + BookHasNoChaptersError: If the book has no indexed chapters. The deletion is rolled + back, so the book's existing phrases stay intact. + + The deletion and regeneration share the caller's session, so they commit together; a + regeneration failure rolls the deletion back. """ started_at = perf_counter() logger.info( @@ -204,37 +146,14 @@ async def recalculate_candidate_phrases_for_book( source.id, source.title, ) - try: - deleted = await delete_phrase_data_for_book(session, source.id) - chapters = await load_book_chapter_texts(session, source.id) - if not chapters: - logger.warning("ebook_candidate_phrase_recalculation_book_empty source_id=%s", source.id) - await session.commit() - return PhraseRecalculationResult( - book_id=source.id, - deleted_candidates=deleted.deleted_candidates, - deleted_protected_phrases=deleted.deleted_protected_phrases, - deleted_aliases=deleted.deleted_aliases, - deleted_mentions=deleted.deleted_mentions, - candidate_phrases=0, - ) - - candidate_count = await generate_candidate_phrases_for_book( - session, - source.id, - series_id=None, - chapters=chapters, - config=config, - nlp=nlp, - metadata=metadata_for_source(source), - replace_all=True, - use_process_pool=use_process_pool, - ) - await session.commit() - except Exception: - await session.rollback() - logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id) - raise + deleted = await delete_phrase_data_for_book(session, source.id) + candidate_count = await generate_candidate_phrases_for_book( + session, + source.id, + series_id=None, + config=config, + replace_all=True, + ) result = PhraseRecalculationResult( book_id=source.id, @@ -258,57 +177,96 @@ async def recalculate_candidate_phrases_for_book( return result -async def generate_candidate_phrases_for_book( - session: AsyncSession, +def generate_candidate_phrases_for_book_in_worker( book_id: int, series_id: int | None, - chapters: Sequence[str], config: EbookSearchConfig, - *, - nlp: SpacyLanguage | None = None, - metadata: Mapping[str, object] | None = None, - replace_all: bool = False, - use_process_pool: bool = False, ) -> int: - """Extract and store candidate phrases for one book without LLM judging. + """Run one book's candidate generation in a pooled worker process. + + The worker has no engine or session to inherit (neither can cross process boundaries), so + it creates its own engine from environment variables, opens the book's session on it, and + disposes the engine once the book is stored. Args: - session (Session): Active database session. book_id (int): Book the candidates belong to. series_id (int | None): Series scope for the stored candidates. - chapters (Sequence[str]): Chapter-like text blocks used for extraction and frequency counts. config (EbookSearchConfig): Runtime phrase-tuning settings. - nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. - metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source. - replace_all (bool): When the caller has already cleared this book's candidates (e.g. a - recalculation), skip the per-candidate existence lookup and bulk-insert new rows. - use_process_pool (bool): Run the CPU-bound extraction in a worker process to avoid - serializing concurrent requests behind the GIL. Ignored when ``nlp`` is set, since - the spaCy pipeline cannot be sent to a worker process. Returns: int: Number of candidate phrase rows stored. """ + + async def generate_with_worker_engine() -> int: + engine = get_async_postgres_engine(name="RICHIE", vector_engine=True, pool_size=1) + try: + async with AsyncSession(engine, expire_on_commit=False) as session: + return await generate_candidate_phrases_for_book( + session, + book_id, + series_id, + config, + ) + finally: + await engine.dispose() + + return asyncio.run(generate_with_worker_engine()) + + +async def generate_candidate_phrases_for_book( + session: AsyncSession, + book_id: int, + series_id: int | None, + config: EbookSearchConfig, + *, + replace_all: bool = False, +) -> int: + """Load a book's chapters and metadata, extract candidate phrases, and store them without LLM judging. + + The session commits only when the whole book succeeds; any failure rolls the session back, + which also restores rows the caller deleted in the same transaction (e.g. a recalculation). + + Args: + session (AsyncSession): Active database session; committed on success, rolled back on failure. + book_id (int): Book the candidates belong to. + series_id (int | None): Series scope for the stored candidates. + config (EbookSearchConfig): Runtime phrase-tuning settings. + replace_all (bool): When the caller has already cleared this book's candidates (e.g. a + recalculation), skip the per-candidate existence lookup and bulk-insert new rows. + + Returns: + int: Number of candidate phrase rows stored. + + Raises: + BookHasNoChaptersError: If the book has no indexed chapter text. + """ started_at = perf_counter() - book_text = "\n\n".join(chapters) - if use_process_pool and nlp is None: - limited_candidates = await extract_phrase_candidates_in_pool(book_text, chapters, config, metadata=metadata) - else: - limited_candidates = extract_phrase_candidates_for_book( + chapters = await load_book_chapter_texts(session, book_id) + if not chapters: + await session.rollback() + message = f"book {book_id} has no indexed chapters" + raise BookHasNoChaptersError(message) + metadata = await metadata_for_source_id(session, book_id) + try: + book_text = "\n\n".join(chapters) + candidates = extract_phrase_candidates_for_book( book_text, chapters, config, - nlp=nlp, metadata=metadata, ) - saved_count = await store_candidate_phrases_for_book( - session, - book_id, - series_id, - limited_candidates, - config, - replace_all=replace_all, - ) + saved_count = await store_candidate_phrases_for_book( + session, + book_id, + series_id, + candidates, + config, + replace_all=replace_all, + ) + await session.commit() + except Exception: + await session.rollback() + raise logger.info( "ebook_candidate_phrase_generation_book_duration book_id=%s candidates=%s duration_ms=%.1f", book_id, diff --git a/python/ebook_search/protected_phrases/pool.py b/python/ebook_search/protected_phrases/pool.py index 68fdec4..7104797 100644 --- a/python/ebook_search/protected_phrases/pool.py +++ b/python/ebook_search/protected_phrases/pool.py @@ -9,21 +9,11 @@ or server threads. from __future__ import annotations -import asyncio import logging import multiprocessing import os from concurrent.futures import ProcessPoolExecutor from threading import Lock -from typing import TYPE_CHECKING - -from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book - -if TYPE_CHECKING: - from collections.abc import Mapping, Sequence - - from python.ebook_search.config import EbookSearchConfig - from python.ebook_search.protected_phrases.models import PhraseCandidate logger = logging.getLogger(__name__) @@ -66,36 +56,3 @@ def shutdown_extraction_pool() -> None: _extraction_pool.pool.shutdown(wait=False, cancel_futures=True) _extraction_pool.pool = None logger.info("ebook_phrase_extraction_pool_shutdown") - - -async def extract_phrase_candidates_in_pool( - book_text: str, - chapters: Sequence[str], - config: EbookSearchConfig, - *, - metadata: Mapping[str, object] | None, -) -> list[PhraseCandidate]: - """Run book phrase extraction in a worker process and await the result. - - Only the CPU-bound extraction runs in the worker; the caller keeps all database work in the - request process. The spaCy pipeline is not supported here because it is not picklable, so - this always runs the non-spaCy extraction path. - - Args: - book_text (str): Full book text used for extraction. - chapters (Sequence[str]): Chapter-like text blocks used for frequency counts. - config (EbookSearchConfig): Runtime phrase-tuning settings. - metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source. - - Returns: - list[PhraseCandidate]: Scored candidates sorted best-first and capped per book. - """ - pool = get_extraction_pool(config.protected_phrase_extraction_workers) - future = pool.submit( - extract_phrase_candidates_for_book, - book_text, - list(chapters), - config, - metadata=dict(metadata) if metadata is not None else None, - ) - return await asyncio.wrap_future(future) diff --git a/tests/ebook_search/test_protected_phrases.py b/tests/ebook_search/test_protected_phrases.py index 8444142..33b0803 100644 --- a/tests/ebook_search/test_protected_phrases.py +++ b/tests/ebook_search/test_protected_phrases.py @@ -2,15 +2,16 @@ from __future__ import annotations +from concurrent.futures import ThreadPoolExecutor from datetime import UTC, datetime from typing import TYPE_CHECKING import pytest from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine -from sqlalchemy.pool import StaticPool from python.ebook_search.config import EbookSearchConfig +from python.ebook_search.protected_phrases import generate_ngrams from python.ebook_search.protected_phrases.config import ( get_bad_ends, get_most_common_words, @@ -55,25 +56,42 @@ from python.orm.richie import ( ) if TYPE_CHECKING: - from collections.abc import AsyncGenerator + from collections.abc import AsyncGenerator, Generator + from pathlib import Path from pytest_mock import MockerFixture @pytest.fixture -async def engine() -> AsyncGenerator[AsyncEngine]: - """Create a shared in-memory async database engine for phrase tests.""" - test_engine = create_async_engine( - "sqlite+aiosqlite:///:memory:", - connect_args={"check_same_thread": False}, - poolclass=StaticPool, - ) +async def engine(tmp_path: Path) -> AsyncGenerator[AsyncEngine]: + """Create a file-backed async database engine that worker threads can also reach.""" + test_engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'phrases.db'}") async with test_engine.begin() as connection: await connection.run_sync(RichieBase.metadata.create_all) yield test_engine await test_engine.dispose() +@pytest.fixture +def worker_pool(engine: AsyncEngine, mocker: MockerFixture) -> Generator[ThreadPoolExecutor]: + """Run pooled candidate generation in threads against the test database. + + Spawned worker processes can see neither the test database nor test patches, so the shared + extraction pool is replaced with a thread pool and worker engines are built for the test + database instead of from Postgres environment variables. + """ + thread_pool = ThreadPoolExecutor(max_workers=1) + database_url = engine.url.render_as_string(hide_password=False) + mocker.patch.object(generate_ngrams, "get_extraction_pool", return_value=thread_pool) + mocker.patch.object( + generate_ngrams, + "get_async_postgres_engine", + side_effect=lambda **_kwargs: create_async_engine(database_url), + ) + yield thread_pool + thread_pool.shutdown(wait=True) + + @pytest.fixture async def session(engine: AsyncEngine) -> AsyncGenerator[AsyncSession]: """Provide a session on the shared in-memory database.""" @@ -271,7 +289,9 @@ async def test_index_chunk_phrase_mentions_uses_normalized_window_lookup( } +@pytest.mark.usefixtures("worker_pool") async def test_generate_candidate_phrases_for_books_stores_unjudged_candidates( + engine: AsyncEngine, session: AsyncSession, config: EbookSearchConfig, ) -> None: @@ -300,8 +320,7 @@ async def test_generate_candidate_phrases_for_books_stores_unjudged_candidates( } ) - result = await generate_candidate_phrases_for_books(session, build_config) - await session.commit() + result = await generate_candidate_phrases_for_books(engine, build_config) candidate = await session.scalar(select(EbookCandidatePhrase)) assert result.books_seen == 1 @@ -313,7 +332,9 @@ async def test_generate_candidate_phrases_for_books_stores_unjudged_candidates( assert await session.scalar(select(EbookProtectedPhrase)) is None +@pytest.mark.usefixtures("worker_pool") async def test_generate_candidate_phrases_for_books_filters_one_token_and_one_use_candidates( + engine: AsyncEngine, session: AsyncSession, config: EbookSearchConfig, ) -> None: @@ -377,8 +398,7 @@ async def test_generate_candidate_phrases_for_books_filters_one_token_and_one_us } ) - result = await generate_candidate_phrases_for_books(session, build_config) - await session.commit() + result = await generate_candidate_phrases_for_books(engine, build_config) candidates = list(await session.scalars(select(EbookCandidatePhrase))) phrase_norms = {candidate.phrase_norm for candidate in candidates} @@ -394,7 +414,9 @@ async def test_generate_candidate_phrases_for_books_filters_one_token_and_one_us assert all(not all(token in common_words for token in candidate.phrase_norm.split()) for candidate in candidates) +@pytest.mark.usefixtures("worker_pool") async def test_generate_candidate_phrases_for_books_commits_after_each_book( + engine: AsyncEngine, session: AsyncSession, config: EbookSearchConfig, mocker: MockerFixture, @@ -431,7 +453,7 @@ async def test_generate_candidate_phrases_for_books_commits_after_each_book( ] ) await session.commit() - commit_spy = mocker.spy(session, "commit") + commit_spy = mocker.spy(AsyncSession, "commit") build_config = config.model_copy( update={ "protected_phrase_max_candidates_per_book": 1, @@ -440,13 +462,91 @@ async def test_generate_candidate_phrases_for_books_commits_after_each_book( } ) - result = await generate_candidate_phrases_for_books(session, build_config) + result = await generate_candidate_phrases_for_books(engine, build_config) assert result.books_built == 2 assert result.candidate_phrases == 2 assert commit_spy.call_count == 2 +@pytest.mark.usefixtures("worker_pool") +async def test_generate_candidate_phrases_for_books_failure_keeps_committed_books( + engine: AsyncEngine, + session: AsyncSession, + config: EbookSearchConfig, + mocker: MockerFixture, +) -> None: + """A failing book should be logged and skipped while the other books stay committed.""" + first = await add_source(session) + second = await add_source(session, file_path="/library/book-2.epub", file_sha256="z" * 64) + session.add_all( + [ + EbookChunk( + id=1, + source_id=first.id, + chapter_id=None, + chunk_index=0, + text="lock in lock in lock in", + token_start=0, + token_count=6, + page_label=None, + content_sha256="f" * 64, + search_text="lock in lock in lock in", + ), + EbookChunk( + id=2, + source_id=second.id, + chapter_id=None, + chunk_index=0, + text="mage king mage king mage king", + token_start=0, + token_count=6, + page_label=None, + content_sha256="g" * 64, + search_text="mage king mage king mage king", + ), + ] + ) + await session.commit() + build_config = config.model_copy( + update={ + "protected_phrase_max_candidates_per_book": 1, + "phrase_min_tokens": 2, + "phrase_max_tokens": 2, + } + ) + real_store = generate_ngrams.store_candidate_phrases_for_book + first_book_id = first.id + second_book_id = second.id + + async def store_failing_second_book( + store_session: AsyncSession, + book_id: int, + series_id: int | None, + limited_candidates: list[PhraseCandidate], + store_config: EbookSearchConfig, + *, + replace_all: bool = False, + ) -> int: + if book_id == second_book_id: + message = "storage exploded" + raise RuntimeError(message) + return await real_store( + store_session, book_id, series_id, limited_candidates, store_config, replace_all=replace_all + ) + + mocker.patch.object(generate_ngrams, "store_candidate_phrases_for_book", side_effect=store_failing_second_book) + + result = await generate_candidate_phrases_for_books(engine, build_config) + + stored_book_ids = set((await session.scalars(select(EbookCandidatePhrase.book_id))).all()) + assert stored_book_ids == {first_book_id} + assert result.books_seen == 2 + assert result.books_built == 1 + assert result.candidate_phrases == 1 + + +@pytest.mark.usefixtures("worker_pool") async def test_judge_candidate_phrases_for_books_promotes_stored_candidates( engine: AsyncEngine, session: AsyncSession, @@ -480,8 +580,7 @@ async def test_judge_candidate_phrases_for_books_promotes_stored_candidates( "phrase_judge_phrase_workers": 1, } ) - await generate_candidate_phrases_for_books(session, build_config) - await session.commit() + await generate_candidate_phrases_for_books(engine, build_config) mocker.patch( "python.ebook_search.protected_phrases.judge_ngrams.judge_candidate_async", return_value=LLMJudgment( @@ -513,6 +612,7 @@ async def test_judge_candidate_phrases_for_books_promotes_stored_candidates( assert mention.phrase_id == phrase.id +@pytest.mark.usefixtures("worker_pool") async def test_judge_candidate_phrases_for_books_logs_and_continues_after_book_failure( engine: AsyncEngine, session: AsyncSession, @@ -561,8 +661,7 @@ async def test_judge_candidate_phrases_for_books_logs_and_continues_after_book_f "phrase_judge_phrase_workers": 1, } ) - await generate_candidate_phrases_for_books(session, build_config) - await session.commit() + await generate_candidate_phrases_for_books(engine, build_config) def judge_or_fail(_client: object, _config: EbookSearchConfig, candidate: PhraseCandidate) -> LLMJudgment: if candidate.phrase_norm == "lock in": @@ -805,6 +904,7 @@ async def test_recalculate_candidate_phrases_for_book_removes_old_phrase_data( result = await recalculate_candidate_phrases_for_book(session, source, build_config) + session.expire_all() candidates = list(await session.scalars(select(EbookCandidatePhrase))) assert result.deleted_candidates == 1 assert result.deleted_protected_phrases == 1 @@ -817,6 +917,38 @@ async def test_recalculate_candidate_phrases_for_book_removes_old_phrase_data( assert await session.scalar(select(EbookChunkPhraseMention)) is None +async def test_recalculate_candidate_phrases_for_book_aborts_without_chapters( + session: AsyncSession, + config: EbookSearchConfig, +) -> None: + """Recalculating a book with no indexed chapters should raise and leave phrase data intact.""" + source = await add_source(session) + existing_candidate = EbookCandidatePhrase( + book_id=source.id, + series_id=None, + phrase_text="old phrase", + phrase_norm="old phrase", + token_count=2, + source_raw_ngram=True, + raw_count=1, + chapter_count=1, + candidate_score=1.0, + llm_judged=False, + ) + session.add(existing_candidate) + await session.flush() + existing_phrase = await add_phrase(session, source.id, phrase_text="old phrase", phrase_norm="old phrase") + await session.commit() + existing_candidate_id = existing_candidate.id + existing_phrase_id = existing_phrase.id + + with pytest.raises(ValueError, match="no indexed chapters"): + await recalculate_candidate_phrases_for_book(session, source, config) + + assert await session.scalar(select(EbookCandidatePhrase.id)) == existing_candidate_id + assert await session.scalar(select(EbookProtectedPhrase.id)) == existing_phrase_id + + async def test_corpus_phrase_stats_counts_phrases_and_book_coverage(session: AsyncSession) -> None: """Corpus stats should count phrases plus how many books are generated and fully judged.""" unjudged_book = await add_source(session) diff --git a/tests/ebook_search/test_ui.py b/tests/ebook_search/test_ui.py index 89db115..5ce355c 100644 --- a/tests/ebook_search/test_ui.py +++ b/tests/ebook_search/test_ui.py @@ -567,33 +567,8 @@ def test_admin_page_shows_protected_phrase_stats(mocker: MockerFixture) -> None: assert value in response.text -def test_ui_add_missing_phrases_generates_only_missing_books(mocker: MockerFixture) -> None: - captured: dict[str, object] = {} - - def fake_generate(_session, _config, *, only_missing=False): - captured["only_missing"] = only_missing - return PhraseCandidateGenerationResult(books_seen=3, books_built=2, candidate_phrases=42) - - mocker.patch( - "python.ebook_search.api.routes.admin.generate_candidate_phrases_for_books", - side_effect=fake_generate, - ) - patch_app_runtime(mocker) - app = create_app() - - with TestClient(app) as client: - response = client.post("/admin/phrases/generate-missing") - - assert response.status_code == 200 - assert captured["only_missing"] is True - assert "42 candidates stored" in response.text - - def test_ui_regenerate_all_phrases_generates_every_book(mocker: MockerFixture) -> None: - captured: dict[str, object] = {} - - def fake_generate(_session, _config, *, only_missing=False): - captured["only_missing"] = only_missing + def fake_generate(_session, _config): return PhraseCandidateGenerationResult(books_seen=5, books_built=5, candidate_phrases=99) mocker.patch( @@ -607,7 +582,6 @@ def test_ui_regenerate_all_phrases_generates_every_book(mocker: MockerFixture) - response = client.post("/admin/phrases/generate-all") assert response.status_code == 200 - assert captured["only_missing"] is False assert "5 of 5 books" in response.text