Richie
c7cd63f8e4
Convert the ebook-search web app to async end to end and add concurrency to the protected-phrase extraction and judging pipeline so large books no longer block the event loop or the UI. ORM / infra: - Add get_async_postgres_engine and factor shared URL/connect_args building into build_postgres_url (reused by the sync and async engine builders) - Add async FastAPI session helpers (get_async_db, AsyncDbSession) with expire_on_commit=False to avoid implicit IO under asyncio App: - Use AsyncEngine/AsyncSession throughout routes, search, ingest, embeddings, answer, rerank and LLM calls; convert handlers to async - Share a single httpx.AsyncClient in app state for LLM requests; size the connection pool for concurrent phrase-judging workers - Add judge_tasks: run per-book judging as tracked background tasks so a book already being judged isn't double-queued Protected phrases: - Add a process pool (pool.py) and worker-count config (extraction/judge book/phrase workers) to parallelize candidate generation and judging - Split admin actions into all/missing variants for generation and judging Config: - Add protected_phrase_extraction_workers, phrase_judge_book_workers, phrase_judge_phrase_workers
371 lines
14 KiB
Python
371 lines
14 KiB
Python
"""Book-level orchestration for candidate n-gram generation and recalculation."""
|
|
|
|
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 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 extract_phrase_candidates_in_pool, 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,
|
|
new_candidate_row,
|
|
prune_unstorable_unjudged_candidate_phrases,
|
|
)
|
|
from python.orm.richie import EbookCandidatePhrase, EbookSource
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Mapping, Sequence
|
|
from concurrent.futures import Future
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
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__)
|
|
|
|
|
|
async def generate_candidate_phrases_for_books(
|
|
session: AsyncSession,
|
|
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.
|
|
|
|
Args:
|
|
session (Session): Active database session.
|
|
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.
|
|
"""
|
|
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)
|
|
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,
|
|
)
|
|
|
|
outcomes = await generate_candidates_for_sources_pooled(session, sources, config)
|
|
|
|
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 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.
|
|
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.
|
|
"""
|
|
started_at = perf_counter()
|
|
logger.info(
|
|
"ebook_candidate_phrase_recalculation_start source_id=%s title=%r",
|
|
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
|
|
|
|
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
|
|
|
|
|
|
async def generate_candidate_phrases_for_book(
|
|
session: AsyncSession,
|
|
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.
|
|
|
|
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.
|
|
"""
|
|
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(
|
|
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,
|
|
)
|
|
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
|