Richie
c9f859cac5
Move protected phrase detection into the retrieval gather so it runs concurrently with vector and BM25 candidates instead of sequentially before them. Make the search API accept real bool form fields for rerank/phrase_matching, gate phrase matching on both the request and config kill switch, and reflow log f-strings for readability.
296 lines
11 KiB
Python
296 lines
11 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(
|
|
f"ebook_candidate_phrase_generation_start {books_seen=} {config.phrase_min_tokens=} "
|
|
f"{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(f"ebook_candidate_phrase_generation_book_failed {source_id=}")
|
|
outcomes.append(BookCandidateResult())
|
|
continue
|
|
saved_count = wrapped_future.result()
|
|
logger.info(f"ebook_candidate_phrase_generation_book_committed {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(
|
|
f"ebook_candidate_phrase_generation_complete {result.books_seen=} {result.books_built=} "
|
|
f"{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(f"ebook_candidate_phrase_recalculation_start {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(
|
|
f"ebook_candidate_phrase_recalculation_complete {source.id=} {result.deleted_candidates=} "
|
|
f"{result.deleted_protected_phrases=} {result.deleted_aliases=} {result.deleted_mentions=} "
|
|
f"{result.candidate_phrases=} duration_ms={(perf_counter() - started_at) * 1000:.1f}"
|
|
)
|
|
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(
|
|
f"ebook_candidate_phrase_generation_book_duration {book_id=} {saved_count=} "
|
|
f"duration_ms={(perf_counter() - started_at) * 1000:.1f}"
|
|
)
|
|
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(
|
|
f"ebook_candidate_phrase_save_start {book_id=} candidates={len(limited_candidates)} mode=bulk_insert"
|
|
)
|
|
else:
|
|
pruned_count = await prune_unstorable_unjudged_candidate_phrases(session, book_id, config)
|
|
logger.info(
|
|
f"ebook_candidate_phrase_save_start {book_id=} candidates={len(limited_candidates)} {pruned_count=}"
|
|
)
|
|
saved_count = await bulk_upsert_unjudged_candidates(session, book_id, series_id, limited_candidates)
|
|
logger.info(
|
|
f"ebook_candidate_phrase_save_complete {book_id=} {saved_count=} "
|
|
f"save_ms={(perf_counter() - save_started_at) * 1000:.1f}"
|
|
)
|
|
return saved_count
|