feat(ebook): migrate to async DB/HTTP and parallelize phrase pipeline

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
This commit is contained in:
2026-07-12 17:51:06 -04:00
parent 63e0b0dd3b
commit bfb3463fd0
29 changed files with 1824 additions and 769 deletions
+244 -72
View File
@@ -11,7 +11,11 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
from python.ebook_search.protected_phrases.extraction import minimum_candidate_raw_count
from python.ebook_search.protected_phrases.models import PhraseCandidate, PhraseRecalculationResult
from python.ebook_search.protected_phrases.models import (
CorpusPhraseStats,
PhraseCandidate,
PhraseRecalculationResult,
)
from python.ebook_search.protected_phrases.text_normalization import normalize_text
from python.orm.richie import (
EbookCandidatePhrase,
@@ -27,7 +31,7 @@ if TYPE_CHECKING:
from sqlalchemy.dialects.postgresql.dml import Insert as PostgresInsert
from sqlalchemy.dialects.sqlite.dml import Insert as SqliteInsert
from sqlalchemy.orm import Session
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import LLMJudgment
@@ -36,14 +40,14 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def dialect_insert(session: Session, table: type[TableBase]) -> PostgresInsert | SqliteInsert:
def dialect_insert(session: AsyncSession, table: type[TableBase]) -> PostgresInsert | SqliteInsert:
"""Return a dialect-specific INSERT construct that supports ``ON CONFLICT DO UPDATE``.
Production runs on PostgreSQL while tests run on SQLite; both support upserts with
compatible SQLAlchemy constructs, so the correct one is chosen from the bound dialect.
Args:
session (Session): Active database session whose bind selects the dialect.
session (AsyncSession): Active database session whose bind selects the dialect.
table (type[TableBase]): Mapped table to insert into.
Returns:
@@ -54,35 +58,33 @@ def dialect_insert(session: Session, table: type[TableBase]) -> PostgresInsert |
return pg_insert(table)
def load_book_text(session: Session, book_id: int) -> str:
async def load_book_text(session: AsyncSession, book_id: int) -> str:
"""Load a book's indexed chunk text as one string for phrase extraction.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose chunk text is loaded.
Returns:
str: The book's chunk text joined into a single string.
"""
texts = session.scalars(
texts = await session.scalars(
select(EbookChunk.text).where(EbookChunk.source_id == book_id).order_by(EbookChunk.chunk_index)
)
return "\n\n".join(stripped for text in texts if (stripped := text.strip()))
def load_book_chapter_texts(session: Session, book_id: int) -> list[str]:
async def load_book_chapter_texts(session: AsyncSession, book_id: int) -> list[str]:
"""Reconstruct chapter-like text blocks from indexed chunks for phrase extraction.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose chunks are grouped into chapters.
Returns:
list[str]: Non-empty chapter-like text blocks in chunk order.
"""
rows = session.execute(
rows = await session.execute(
select(EbookChunk.chapter_id, EbookChunk.text)
.where(EbookChunk.source_id == book_id)
.order_by(EbookChunk.chunk_index)
@@ -127,11 +129,11 @@ def metadata_for_source(source: EbookSource) -> dict[str, object | None]:
}
def metadata_for_source_id(session: Session, source_id: int) -> dict[str, object | None]:
async def metadata_for_source_id(session: AsyncSession, source_id: int) -> dict[str, object | None]:
"""Return phrase extraction metadata for one indexed source by id.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
source_id (int): Id of the indexed source to read metadata from.
Returns:
@@ -140,70 +142,128 @@ def metadata_for_source_id(session: Session, source_id: int) -> dict[str, object
Raises:
ValueError: If no source exists with the given id.
"""
source = session.get(EbookSource, source_id)
source = await session.get(EbookSource, source_id)
if source is None:
msg = f"No indexed source with id {source_id}"
raise ValueError(msg)
return metadata_for_source(source)
def count_protected_phrases(session: Session, book_id: int) -> int:
async def count_protected_phrases(session: AsyncSession, book_id: int) -> int:
"""Count stored protected phrases for one book.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose protected phrases are counted.
Returns:
int: Number of protected phrases stored for the book.
"""
return session.scalars(
select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)
return (
await session.scalars(
select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)
)
).one()
def count_unjudged_candidates(session: Session, book_id: int, config: EbookSearchConfig) -> int:
async def count_unjudged_candidates(session: AsyncSession, book_id: int, config: EbookSearchConfig) -> int:
"""Count storable candidate rows for a book that have not yet been judged.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose unjudged candidates are counted.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
Returns:
int: Number of storable, unjudged candidate rows for the book.
"""
return session.scalars(
select(func.count(EbookCandidatePhrase.id)).where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(False),
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
return (
await session.scalars(
select(func.count(EbookCandidatePhrase.id)).where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(False),
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
)
)
).one()
def load_candidates_for_judgment(
session: Session,
book_id: int,
judgment_limit: int,
config: EbookSearchConfig,
) -> Sequence[EbookCandidatePhrase]:
"""Load the top unjudged candidate rows for a book.
Common-word phrases are filtered before storage (``filter_storable_candidates``), so
no re-check is needed here.
async def corpus_phrase_stats(session: AsyncSession) -> CorpusPhraseStats:
"""Summarize candidate and protected phrase coverage across the whole corpus.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
Returns:
CorpusPhraseStats: Corpus-wide phrase counts and per-book coverage counts.
"""
total_books = (await session.scalars(select(func.count(EbookSource.id)))).one()
candidate_phrases, judged_candidates, books_with_candidates, books_with_unjudged = (
await session.execute(
select(
func.count(EbookCandidatePhrase.id),
func.count(EbookCandidatePhrase.id).filter(EbookCandidatePhrase.llm_judged.is_(True)),
func.count(func.distinct(EbookCandidatePhrase.book_id)),
func.count(func.distinct(EbookCandidatePhrase.book_id)).filter(
EbookCandidatePhrase.llm_judged.is_(False)
),
)
)
).one()
protected_phrases = (await session.scalars(select(func.count(EbookProtectedPhrase.id)))).one()
return CorpusPhraseStats(
total_books=total_books,
books_with_candidates=books_with_candidates,
books_fully_judged=books_with_candidates - books_with_unjudged,
candidate_phrases=candidate_phrases,
judged_candidates=judged_candidates,
unjudged_candidates=candidate_phrases - judged_candidates,
protected_phrases=protected_phrases,
)
async def book_ids_pending_first_judgment(session: AsyncSession) -> list[int]:
"""Return books that have candidate phrases but no judged candidates yet.
Args:
session (AsyncSession): Active database session.
Returns:
list[int]: Book ids with candidates where judging has never run, ordered by id.
"""
judged_books = select(EbookCandidatePhrase.book_id).where(EbookCandidatePhrase.llm_judged.is_(True)).distinct()
return list(
(
await session.scalars(
select(EbookCandidatePhrase.book_id)
.where(EbookCandidatePhrase.book_id.not_in(judged_books))
.distinct()
.order_by(EbookCandidatePhrase.book_id)
)
).all()
)
async def load_candidates_for_judgment(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
) -> Sequence[EbookCandidatePhrase]:
"""Load every storable unjudged candidate row for a book.
Rows may have been stored before the current junk filters and score weights existed, so
callers re-check :func:`is_junk_phrase` and rescore before selecting what to judge.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose candidates are loaded.
judgment_limit (int): Maximum number of candidate rows to return.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
Returns:
Sequence[EbookCandidatePhrase]: Top storable, unjudged candidate rows ordered by score.
Sequence[EbookCandidatePhrase]: Storable, unjudged candidate rows ordered by stored score.
"""
return session.scalars(
query = (
select(EbookCandidatePhrase)
.where(
EbookCandidatePhrase.book_id == book_id,
@@ -216,8 +276,8 @@ def load_candidates_for_judgment(
EbookCandidatePhrase.raw_count.desc(),
EbookCandidatePhrase.id,
)
.limit(judgment_limit)
).all()
)
return (await session.scalars(query)).all()
def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
@@ -248,25 +308,23 @@ def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
)
def save_candidate_to_db(
session: Session,
def candidate_row_values(
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
*,
judgment: LLMJudgment | None,
) -> EbookCandidatePhrase:
"""Insert or update one candidate phrase row.
) -> dict[str, object]:
"""Build the column values for one candidate phrase upsert.
Args:
session (Session): Active database session.
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope stored on the row.
candidate (PhraseCandidate): Candidate whose fields are written to the row.
judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged.
Returns:
EbookCandidatePhrase: The inserted or updated candidate row.
dict[str, object]: Column values keyed by column name.
"""
values: dict[str, object] = {
"book_id": book_id,
@@ -296,6 +354,30 @@ def save_candidate_to_db(
llm_category=judgment.category,
llm_reason=judgment.reason,
)
return values
async def save_candidate_to_db(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
*,
judgment: LLMJudgment | None,
) -> EbookCandidatePhrase:
"""Insert or update one candidate phrase row.
Args:
session (AsyncSession): Active database session.
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope stored on the row.
candidate (PhraseCandidate): Candidate whose fields are written to the row.
judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged.
Returns:
EbookCandidatePhrase: The inserted or updated candidate row.
"""
values = candidate_row_values(book_id, series_id, candidate, judgment=judgment)
# Preserve an existing judgment when this call is only refreshing candidate fields.
skip_update = {"book_id", "phrase_norm"}
@@ -306,11 +388,94 @@ def save_candidate_to_db(
index_elements=["book_id", "phrase_norm"],
set_={column: insert_statement.excluded[column] for column in values if column not in skip_update},
).returning(EbookCandidatePhrase)
return session.scalars(statement, execution_options={"populate_existing": True}).one()
return (await session.scalars(statement, execution_options={"populate_existing": True})).one()
def upsert_protected_phrase(
session: Session,
BULK_CANDIDATE_UPSERT_CHUNK = 1000
async def bulk_upsert_unjudged_candidates(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidates: Sequence[PhraseCandidate],
) -> int:
"""Insert or update many freshly extracted candidate rows in chunked multi-row upserts.
Saving one row per statement costs one database round trip per candidate, which dominated
generation time for full books, so candidates are written ``BULK_CANDIDATE_UPSERT_CHUNK``
rows per statement instead. Existing judgments and sample contexts are never overwritten:
fresh extractions carry no contexts, and ``llm_judged`` plus the ``llm_*`` columns are left
out of the conflict update. Candidates must have unique ``phrase_norm`` values, as produced
by extraction, since one multi-row upsert cannot touch the same row twice.
Args:
session (Session): Active database session.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope stored on the rows.
candidates (Sequence[PhraseCandidate]): Freshly extracted candidates to persist.
Returns:
int: Number of candidate rows written.
"""
values = [
candidate_row_values(book_id, series_id, candidate, judgment=None)
for candidate in candidates
if not candidate.sample_contexts
]
if len(values) != len(candidates):
msg = "bulk_upsert_unjudged_candidates only accepts freshly extracted candidates without sample contexts"
raise ValueError(msg)
skip_update = {"book_id", "phrase_norm", "llm_judged"}
for chunk_start in range(0, len(values), BULK_CANDIDATE_UPSERT_CHUNK):
chunk = values[chunk_start : chunk_start + BULK_CANDIDATE_UPSERT_CHUNK]
insert_statement = dialect_insert(session, EbookCandidatePhrase).values(chunk)
statement = insert_statement.on_conflict_do_update(
index_elements=["book_id", "phrase_norm"],
set_={column: insert_statement.excluded[column] for column in chunk[0] if column not in skip_update},
)
await session.execute(statement)
return len(values)
def new_candidate_row(book_id: int, series_id: int | None, candidate: PhraseCandidate) -> EbookCandidatePhrase:
"""Build a fresh unjudged candidate row without checking for an existing one.
Unlike :func:`save_candidate_to_db`, this does no lookup, so it is only safe when the caller
guarantees there is no existing row for ``(book_id, candidate.phrase_norm)`` — for example
right after :func:`delete_phrase_data_for_book` has cleared the book.
Args:
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope stored on the row.
candidate (PhraseCandidate): Candidate whose fields are written to the row.
Returns:
EbookCandidatePhrase: A new, unattached candidate row.
"""
row = EbookCandidatePhrase(book_id=book_id, phrase_norm=candidate.phrase_norm)
row.llm_judged = False
row.series_id = series_id
row.phrase_text = candidate.phrase_text
row.token_count = candidate.token_count
row.source_raw_ngram = candidate.source_raw_ngram
row.source_yake = candidate.source_yake
row.source_spacy_ner = candidate.source_spacy_ner
row.source_spacy_noun_chunk = candidate.source_spacy_noun_chunk
row.source_capitalized = candidate.source_capitalized
row.source_metadata = candidate.source_metadata
row.spacy_label = candidate.spacy_label
row.raw_count = candidate.raw_count
row.chapter_count = candidate.chapter_count
row.yake_score = candidate.yake_score
row.candidate_score = candidate.candidate_score
if candidate.sample_contexts:
row.sample_contexts = list(candidate.sample_contexts)
return row
async def upsert_protected_phrase(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
@@ -320,7 +485,7 @@ def upsert_protected_phrase(
"""Insert or update one accepted protected phrase and its aliases.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book the protected phrase belongs to.
series_id (int | None): Series scope stored on the phrase.
candidate (PhraseCandidate): Candidate the phrase was promoted from.
@@ -360,18 +525,22 @@ def upsert_protected_phrase(
column: insert_statement.excluded[column] for column in values if column not in {"book_id", "phrase_norm"}
},
).returning(EbookProtectedPhrase)
row = session.scalars(statement, execution_options={"populate_existing": True}).one()
row = (await session.scalars(statement, execution_options={"populate_existing": True})).one()
for alias_text in judgment.aliases:
upsert_phrase_alias(session, row, alias_text)
await upsert_phrase_alias(session, row, alias_text)
return row
def upsert_phrase_alias(session: Session, phrase: EbookProtectedPhrase, alias_text: str) -> EbookPhraseAlias | None:
async def upsert_phrase_alias(
session: AsyncSession,
phrase: EbookProtectedPhrase,
alias_text: str,
) -> EbookPhraseAlias | None:
"""Insert or update one protected phrase alias.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
phrase (EbookProtectedPhrase): Protected phrase the alias points to.
alias_text (str): Alias surface form to store.
@@ -395,7 +564,7 @@ def upsert_phrase_alias(session: Session, phrase: EbookProtectedPhrase, alias_te
"confidence": insert_statement.excluded.confidence,
},
).returning(EbookPhraseAlias)
return session.scalars(statement, execution_options={"populate_existing": True}).one()
return (await session.scalars(statement, execution_options={"populate_existing": True})).one()
def make_canonical_id(judgment: LLMJudgment, phrase_norm: str) -> str:
@@ -426,15 +595,15 @@ def slugify_identifier(value: str) -> str:
return slug.strip("_") or "unknown"
def prune_unstorable_unjudged_candidate_phrases(
session: Session,
async def prune_unstorable_unjudged_candidate_phrases(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
) -> int:
"""Delete old unjudged candidate rows that no longer satisfy storage filters.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose stale candidates are pruned.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
@@ -442,7 +611,7 @@ def prune_unstorable_unjudged_candidate_phrases(
int: Number of candidate rows deleted.
"""
deleted = rowcount(
session.execute(
await session.execute(
delete(EbookCandidatePhrase).where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(False),
@@ -464,40 +633,42 @@ def prune_unstorable_unjudged_candidate_phrases(
return deleted
def delete_phrase_data_for_book(session: Session, book_id: int) -> PhraseRecalculationResult:
async def delete_phrase_data_for_book(session: AsyncSession, book_id: int) -> PhraseRecalculationResult:
"""Delete all candidate, protected, alias, and mention phrase data for one book.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose phrase data is deleted.
Returns:
PhraseRecalculationResult: Deleted-row counts with ``candidate_phrases`` set to 0.
"""
protected_ids = session.scalars(
select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id)
protected_ids = (
await session.scalars(select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id))
).all()
deleted_aliases = 0
if protected_ids:
deleted_aliases = rowcount(
session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids)))
await session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids)))
)
deleted_mentions = rowcount(
session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
)
if protected_ids:
deleted_mentions += rowcount(
session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids)))
await session.execute(
delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids))
)
)
deleted_protected = rowcount(
session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id))
await session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id))
)
deleted_candidates = rowcount(
session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id))
await session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id))
)
session.flush()
await session.flush()
logger.info(
"ebook_candidate_phrase_data_deleted book_id=%s candidates=%s protected=%s aliases=%s mentions=%s",
book_id,
@@ -515,6 +686,7 @@ def delete_phrase_data_for_book(session: Session, book_id: int) -> PhraseRecalcu
candidate_phrases=0,
)
def rowcount(result: object) -> int:
"""Return a safe integer rowcount from a SQLAlchemy execution result.