Richie/dotfiles
1
0
Fork 0
Code Issues 1 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

fix(protected-phrases): isolate phrase generation per book

Browse Source 
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:
Richie
2026-07-24 11:38:50 -04:00
parent 29a51eb1b8
commit 58be234d7f
6 changed files with 300 additions and 306 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ebook_search/api/routes/admin.py
+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",
python/ebook_search/api/routes/page.py
+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",
python/ebook_search/protected_phrases/generate_ngrams.py
+140 -182
View File
@@ -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,
python/ebook_search/protected_phrases/pool.py
-43
View File
@@ -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)