Richie
58be234d7f
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.
329 lines
12 KiB
Python
329 lines
12 KiB
Python
"""Book-level orchestration for candidate n-gram generation and recalculation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
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 (
|
|
BookCandidateResult,
|
|
PhraseCandidateGenerationResult,
|
|
PhraseRecalculationResult,
|
|
)
|
|
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_id,
|
|
new_candidate_row,
|
|
prune_unstorable_unjudged_candidate_phrases,
|
|
)
|
|
from python.orm.common import get_async_postgres_engine
|
|
from python.orm.richie import EbookSource
|
|
|
|
if TYPE_CHECKING:
|
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
|
from python.ebook_search.config import EbookSearchConfig
|
|
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(
|
|
engine: AsyncEngine,
|
|
config: EbookSearchConfig,
|
|
) -> PhraseCandidateGenerationResult:
|
|
"""Create or refresh candidate phrases for indexed books without calling the LLM judge.
|
|
|
|
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:
|
|
engine (AsyncEngine): Engine used to read the book list in this process.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
|
|
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.
|
|
"""
|
|
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",
|
|
books_seen,
|
|
config.phrase_min_tokens,
|
|
config.phrase_max_tokens,
|
|
config.protected_phrase_max_candidates_per_book,
|
|
)
|
|
|
|
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,
|
|
books_built=sum(1 for outcome in outcomes if outcome.built),
|
|
candidate_phrases=sum(outcome.candidates for outcome in outcomes),
|
|
)
|
|
logger.info(
|
|
"ebook_candidate_phrase_generation_complete books_seen=%s books_built=%s candidate_total=%s",
|
|
result.books_seen,
|
|
result.books_built,
|
|
result.candidate_phrases,
|
|
)
|
|
return result
|
|
|
|
|
|
async def recalculate_candidate_phrases_for_book(
|
|
session: AsyncSession,
|
|
source: EbookSource,
|
|
config: EbookSearchConfig,
|
|
) -> PhraseRecalculationResult:
|
|
"""Remove all book phrase data, regenerate candidates, and commit the completed book.
|
|
|
|
Args:
|
|
session (AsyncSession): Active database session; deletion and regeneration commit on it.
|
|
source (EbookSource): Indexed book to recalculate.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
|
|
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(
|
|
"ebook_candidate_phrase_recalculation_start source_id=%s title=%r",
|
|
source.id,
|
|
source.title,
|
|
)
|
|
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,
|
|
deleted_candidates=deleted.deleted_candidates,
|
|
deleted_protected_phrases=deleted.deleted_protected_phrases,
|
|
deleted_aliases=deleted.deleted_aliases,
|
|
deleted_mentions=deleted.deleted_mentions,
|
|
candidate_phrases=candidate_count,
|
|
)
|
|
logger.info(
|
|
"ebook_candidate_phrase_recalculation_complete source_id=%s deleted_candidates=%s "
|
|
"deleted_protected=%s deleted_aliases=%s deleted_mentions=%s candidates=%s duration_ms=%.1f",
|
|
source.id,
|
|
result.deleted_candidates,
|
|
result.deleted_protected_phrases,
|
|
result.deleted_aliases,
|
|
result.deleted_mentions,
|
|
result.candidate_phrases,
|
|
(perf_counter() - started_at) * 1000,
|
|
)
|
|
return result
|
|
|
|
|
|
def generate_candidate_phrases_for_book_in_worker(
|
|
book_id: int,
|
|
series_id: int | None,
|
|
config: EbookSearchConfig,
|
|
) -> int:
|
|
"""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:
|
|
book_id (int): Book the candidates belong to.
|
|
series_id (int | None): Series scope for the stored candidates.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
|
|
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()
|
|
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,
|
|
metadata=metadata,
|
|
)
|
|
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,
|
|
saved_count,
|
|
(perf_counter() - started_at) * 1000,
|
|
)
|
|
return saved_count
|
|
|
|
|
|
async def store_candidate_phrases_for_book(
|
|
session: AsyncSession,
|
|
book_id: int,
|
|
series_id: int | None,
|
|
limited_candidates: list[PhraseCandidate],
|
|
config: EbookSearchConfig,
|
|
*,
|
|
replace_all: bool = False,
|
|
) -> int:
|
|
"""Persist already-extracted candidate phrase rows for one book without committing.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
book_id (int): Book the candidates belong to.
|
|
series_id (int | None): Series scope for the stored candidates.
|
|
limited_candidates (list[PhraseCandidate]): Scored candidates to persist.
|
|
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
|
replace_all (bool): When the caller has already cleared this book's candidates, skip the
|
|
per-candidate existence lookup and bulk-insert new rows.
|
|
|
|
Returns:
|
|
int: Number of candidate phrase rows stored.
|
|
"""
|
|
save_started_at = perf_counter()
|
|
if replace_all:
|
|
rows = [new_candidate_row(book_id, series_id, candidate) for candidate in limited_candidates]
|
|
session.add_all(rows)
|
|
await session.flush()
|
|
saved_count = len(rows)
|
|
logger.info(
|
|
"ebook_candidate_phrase_save_start book_id=%s candidates=%s mode=bulk_insert",
|
|
book_id,
|
|
len(limited_candidates),
|
|
)
|
|
else:
|
|
pruned_count = await prune_unstorable_unjudged_candidate_phrases(session, book_id, config)
|
|
logger.info(
|
|
"ebook_candidate_phrase_save_start book_id=%s candidates=%s pruned_unstorable=%s",
|
|
book_id,
|
|
len(limited_candidates),
|
|
pruned_count,
|
|
)
|
|
saved_count = await bulk_upsert_unjudged_candidates(session, book_id, series_id, limited_candidates)
|
|
logger.info(
|
|
"ebook_candidate_phrase_save_complete book_id=%s candidates=%s save_ms=%.1f",
|
|
book_id,
|
|
saved_count,
|
|
(perf_counter() - save_started_at) * 1000,
|
|
)
|
|
return saved_count
|