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