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-16 13:09:34 -04:00
parent 9bc0d67962
commit 1f5e527cbc
29 changed files with 1824 additions and 769 deletions
+24 -19
View File
@@ -23,7 +23,8 @@ logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlalchemy.orm import Session
import httpx
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.config import EbookSearchConfig
@@ -65,7 +66,11 @@ class EmbeddingModelStats:
return max(self.total_chunks - self.embedded_chunks, 0)
def embed_texts(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]:
async def embed_texts(
client: httpx.AsyncClient,
texts: Sequence[str],
config: EbookSearchConfig,
) -> list[list[float]]:
"""Embed text with the configured vLLM embedding model."""
logger.info(
"ebook_embed_request_start base_url=%s model=%s count=%s",
@@ -73,7 +78,7 @@ def embed_texts(texts: Sequence[str], config: EbookSearchConfig) -> list[list[fl
config.embedding_model,
len(texts),
)
vectors = request_embeddings(texts, config)
vectors = await request_embeddings(client, texts, config)
expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
for vector in vectors:
if len(vector) != expected_dimension:
@@ -88,28 +93,28 @@ def embed_texts(texts: Sequence[str], config: EbookSearchConfig) -> list[list[fl
return vectors
def embed_query(query: str, config: EbookSearchConfig) -> list[float]:
async def embed_query(client: httpx.AsyncClient, query: str, config: EbookSearchConfig) -> list[float]:
"""Embed a search query with the Qwen retrieval instruction."""
instructed_query = f"Instruct: Retrieve relevant passages for the query.\nQuery: {query}"
return embed_texts([instructed_query], config)[0]
return (await embed_texts(client, [instructed_query], config))[0]
def ensure_embedding_models(session: Session) -> None:
async def ensure_embedding_models(session: AsyncSession) -> None:
"""Ensure supported embedding model rows exist."""
for name, dimension in MODEL_DIMENSIONS.items():
existing = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == name))
existing = await session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == name))
if existing is None:
session.add(EbookEmbeddingModel(name=name, dimension=dimension, is_default=name == "qwen3-embedding-0.6b"))
logger.info("ebook_embedding_model_created model=%s dimension=%s", name, dimension)
session.flush()
await session.flush()
def embedding_model_stats(session: Session) -> list[EmbeddingModelStats]:
async def embedding_model_stats(session: AsyncSession) -> list[EmbeddingModelStats]:
"""Return embedding coverage counts for every supported model."""
total_chunks = session.scalar(select(func.count(EbookChunk.id))) or 0
total_chunks = await session.scalar(select(func.count(EbookChunk.id))) or 0
models = {
model.name: model
for model in session.scalars(
for model in await session.scalars(
select(EbookEmbeddingModel)
.where(EbookEmbeddingModel.name.in_(MODEL_DIMENSIONS))
.order_by(EbookEmbeddingModel.name)
@@ -122,7 +127,7 @@ def embedding_model_stats(session: Session) -> list[EmbeddingModelStats]:
embedded_chunks = 0
if model is not None:
table = get_embedding_table(dimension)
embedded_chunks = session.scalar(select(func.count(table.id)).where(table.model_id == model.id)) or 0
embedded_chunks = await session.scalar(select(func.count(table.id)).where(table.model_id == model.id)) or 0
stats.append(
EmbeddingModelStats(
model_name=model_name,
@@ -134,10 +139,10 @@ def embedding_model_stats(session: Session) -> list[EmbeddingModelStats]:
return stats
def embed_missing_chunks(session: Session, config: EbookSearchConfig) -> int:
async def embed_missing_chunks(session: AsyncSession, client: httpx.AsyncClient, config: EbookSearchConfig) -> int:
"""Embed chunks missing embeddings for the configured model."""
ensure_embedding_models(session)
model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
await ensure_embedding_models(session)
model = await session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
if model is None:
supported_models = ", ".join(MODEL_DIMENSIONS)
msg = f"Unknown embedding model: {config.embedding_model}. Supported models: {supported_models}"
@@ -145,7 +150,7 @@ def embed_missing_chunks(session: Session, config: EbookSearchConfig) -> int:
table = get_embedding_table(model.dimension)
chunks = list(
session.scalars(
await session.scalars(
select(EbookChunk)
.outerjoin(table, (table.chunk_id == EbookChunk.id) & (table.model_id == model.id))
.where(table.id.is_(None))
@@ -158,13 +163,13 @@ def embed_missing_chunks(session: Session, config: EbookSearchConfig) -> int:
return 0
logger.info("ebook_embed_missing_batch_start model=%s count=%s", config.embedding_model, len(chunks))
vectors = embed_texts([chunk.text for chunk in chunks], config)
vectors = await embed_texts(client, [chunk.text for chunk in chunks], config)
rows = [
{"chunk_id": chunk.id, "model_id": model.id, "embedding": vector}
for chunk, vector in zip(chunks, vectors, strict=True)
]
statement = insert(table).values(rows).on_conflict_do_nothing(index_elements=["chunk_id", "model_id"])
session.execute(statement)
session.flush()
await session.execute(statement)
await session.flush()
logger.info("ebook_embed_missing_batch_complete model=%s count=%s", config.embedding_model, len(rows))
return len(rows)