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:
@@ -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()}
|
||||
|
||||
Reference in New Issue
Block a user