fix(protected-phrases): isolate phrase generation per book
treefmt / nix fmt (pull_request) Failing after 5s
pytest / pytest (pull_request) Successful in 28s
test ebook search / test-ebook-search (pull_request) Failing after 35s
build_systems / build-bob (pull_request) Successful in 49s
build_systems / build-brain (pull_request) Successful in 49s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m3s
build_systems / build-jeeves (pull_request) Successful in 2m24s

Run full-book candidate generation inside worker-owned sessions so each book commits independently during backfills. Abort recalculation when a book has no indexed chapters to preserve existing phrase data, and update admin/UI tests for the new generation flow.
This commit is contained in:
2026-07-12 17:51:38 -04:00
parent 547493910e
commit 40f707798e
6 changed files with 300 additions and 306 deletions
+4 -34
View File
@@ -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",
+4 -1
View File
@@ -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",
@@ -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,
@@ -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)
+151 -19
View File
@@ -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)
+1 -27
View File
@@ -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