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-24 11:38:50 -04:00
parent 38c01ec121
commit 2706c4417d
29 changed files with 1824 additions and 769 deletions
@@ -26,7 +26,7 @@ from python.orm.richie import (
if TYPE_CHECKING:
from collections.abc import Iterator, Sequence
from sqlalchemy.orm import Session
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.text_normalization import NormalizedToken
@@ -34,8 +34,8 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def load_phrase_lookup(
session: Session,
async def load_phrase_lookup(
session: AsyncSession,
config: EbookSearchConfig,
*,
book_id: int | None = None,
@@ -44,7 +44,7 @@ def load_phrase_lookup(
"""Load protected phrases and aliases into RAM lookup maps.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
config (EbookSearchConfig): Runtime phrase-tuning settings.
book_id (int | None): Optional book scope to restrict loaded phrases.
series_id (int | None): Optional series scope to restrict loaded phrases.
@@ -65,7 +65,7 @@ def load_phrase_lookup(
if scope_filter is not None:
phrase_statement = phrase_statement.where(scope_filter)
for row in session.execute(phrase_statement):
for row in await session.execute(phrase_statement):
phrase_id = int(row.id)
phrase_norm = str(row.phrase_norm)
norm_to_ids[phrase_norm].append(phrase_id)
@@ -78,7 +78,7 @@ def load_phrase_lookup(
if scope_filter is not None:
alias_statement = alias_statement.where(scope_filter)
for row in session.execute(alias_statement):
for row in await session.execute(alias_statement):
alias_norm = str(row.alias_norm)
alias_to_ids[alias_norm].append(int(row.phrase_id))
max_tokens = max(max_tokens, len(alias_norm.split()))
@@ -197,11 +197,11 @@ def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], loo
return matches
def hydrate_matches(session: Session, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]:
async def hydrate_matches(session: AsyncSession, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]:
"""Fetch protected phrase metadata for raw phrase matches.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
matches (Sequence[PhraseMatch]): Unhydrated matches to enrich.
Returns:
@@ -216,7 +216,7 @@ def hydrate_matches(session: Session, matches: Sequence[PhraseMatch]) -> list[Hy
rows = {
row.id: row
for row in session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids)))
for row in await session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids)))
}
hydrated: list[HydratedPhraseMatch] = []
for match in matches:
@@ -332,8 +332,8 @@ def resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhr
return kept
def detect_protected_phrases_for_query(
session: Session,
async def detect_protected_phrases_for_query(
session: AsyncSession,
query_text: str,
config: EbookSearchConfig,
*,
@@ -344,7 +344,7 @@ def detect_protected_phrases_for_query(
"""Run the full online protected-phrase query-detection pipeline.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
query_text (str): User query text to detect phrases in.
config (EbookSearchConfig): Runtime phrase-tuning settings.
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
@@ -355,13 +355,15 @@ def detect_protected_phrases_for_query(
list[HydratedPhraseMatch]: Hydrated, overlap-resolved phrase matches for the query.
"""
active_lookup = (
lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
lookup
if lookup is not None
else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
)
return resolve_overlaps(hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup)))
return resolve_overlaps(await hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup)))
def index_chunk_phrase_mentions_for_book(
session: Session,
async def index_chunk_phrase_mentions_for_book(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
*,
@@ -371,7 +373,7 @@ def index_chunk_phrase_mentions_for_book(
"""Rebuild chunk phrase mentions for all chunks in one book.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose chunk mentions are rebuilt.
config (EbookSearchConfig): Runtime phrase-tuning settings.
series_id (int | None): Optional series scope for lookup loading.
@@ -381,23 +383,25 @@ def index_chunk_phrase_mentions_for_book(
int: Total number of chunk phrase mentions indexed for the book.
"""
active_lookup = (
lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
lookup
if lookup is not None
else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
)
session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
chunks = session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id))
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
chunks = await session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id))
count = 0
for chunk in chunks:
count += index_chunk_phrase_mentions(session, chunk, lookup=active_lookup)
session.flush()
count += await index_chunk_phrase_mentions(session, chunk, lookup=active_lookup)
await session.flush()
logger.info("ebook_chunk_phrase_mentions_indexed book_id=%s mentions=%s", book_id, count)
return count
def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
async def index_chunk_phrase_mentions(session: AsyncSession, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
"""Store protected phrase mentions for one chunk.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
chunk (EbookChunk): Chunk whose text is scanned for phrase mentions.
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
@@ -405,7 +409,7 @@ def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup:
int: Number of phrase mentions stored for the chunk.
"""
raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup)
hydrated = resolve_overlaps(hydrate_matches(session, raw_matches))
hydrated = resolve_overlaps(await hydrate_matches(session, raw_matches))
for match in hydrated:
session.add(
EbookChunkPhraseMention(
@@ -420,8 +424,8 @@ def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup:
return len(hydrated)
def phrase_hits_for_chunks(
session: Session,
async def phrase_hits_for_chunks(
session: AsyncSession,
*,
chunk_ids: Sequence[int],
phrase_ids: Sequence[int],
@@ -429,7 +433,7 @@ def phrase_hits_for_chunks(
"""Return matched protected phrases with mention counts by chunk id using indexed chunk mentions.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
chunk_ids (Sequence[int]): Chunk ids to look up mentions for.
phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to.
@@ -456,7 +460,7 @@ def phrase_hits_for_chunks(
.order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text)
)
hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list)
for row in session.execute(statement):
for row in await session.execute(statement):
hits[int(row.chunk_id)].append(
ChunkPhraseHit(
phrase_id=int(row.phrase_id),
@@ -467,8 +471,8 @@ def phrase_hits_for_chunks(
return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()}
def phrase_hit_counts_for_chunks(
session: Session,
async def phrase_hit_counts_for_chunks(
session: AsyncSession,
*,
chunk_ids: Sequence[int],
phrase_ids: Sequence[int],
@@ -476,12 +480,12 @@ def phrase_hit_counts_for_chunks(
"""Return phrase-hit counts by chunk id using indexed chunk mentions.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
chunk_ids (Sequence[int]): Chunk ids to count mentions for.
phrase_ids (Sequence[int]): Protected phrase ids to restrict the counts to.
Returns:
dict[int, int]: Total mention count per chunk id.
"""
hits = phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
return {chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in hits.items()}