Richie/dotfiles
1
0
Fork 0
Code Issues 1 Pull Requests 3 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
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:
Richie
2026-07-24 11:38:50 -04:00
parent 38c01ec121
commit 2706c4417d
29 changed files with 1824 additions and 769 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ebook_search/search.py
+54 -54
View File
@@ -2,17 +2,17 @@
from __future__ import annotations
import asyncio
import logging
import re
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING
from pgvector.sqlalchemy import Vector
from sqlalchemy import literal, select
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import Session
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.bm25_corpus import (
BM25CorpusUnavailableError,
@@ -25,7 +25,7 @@ from python.ebook_search.protected_phrases.matching import (
phrase_hits_for_chunks,
)
from python.ebook_search.rerank import rerank_chunks
from python.ebook_search.timing import RuntimeStep, timed_result
from python.ebook_search.timing import RuntimeStep, async_timed_result, timed_result
from python.orm.richie import (
EbookChapter,
EbookChunk,
@@ -36,7 +36,8 @@ from python.orm.richie import (
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from sqlalchemy.engine import Engine
import httpx
from sqlalchemy.ext.asyncio import AsyncEngine
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import HydratedPhraseMatch
@@ -90,8 +91,9 @@ class RetrievalResponse:
timings: tuple[RuntimeStep, ...]
def search_ebooks(
engine: Engine,
async def search_ebooks(
engine: AsyncEngine,
client: httpx.AsyncClient,
query: str,
config: EbookSearchConfig,
*,
@@ -112,16 +114,15 @@ def search_ebooks(
)
timings: list[RuntimeStep] = []
if phrase_matching_enabled:
phrase_matches, timing = timed_result("Protected phrase detection", query_phrase_matches, engine, query, config)
phrase_matches, timing = await async_timed_result(
"Protected phrase detection", query_phrase_matches(engine, query, config)
)
else:
phrase_matches, timing = timed_result("Protected phrase detection skipped", skip_phrase_matches)
timings.append(timing)
retrieval, timing = timed_result(
retrieval, timing = await async_timed_result(
"Hybrid retrieval",
parallel_retrieval,
engine,
query,
config,
parallel_retrieval(engine, client, query, config),
)
timings.extend(retrieval.timings)
timings.append(timing)
@@ -134,19 +135,15 @@ def search_ebooks(
)
timings.append(timing)
if phrase_matching_enabled:
fused, timing = timed_result(
fused, timing = await async_timed_result(
"Phrase mention boost",
apply_phrase_mention_boosts,
engine,
fused,
phrase_matches,
config.phrase_hit_boost,
apply_phrase_mention_boosts(engine, fused, phrase_matches, config.phrase_hit_boost),
)
else:
fused, timing = timed_result("Phrase mention boost skipped", skip_phrase_mention_boosts, fused)
timings.append(timing)
if config.rerank.enabled and rerank:
response, timing = timed_result("Rerank", apply_rerank, query, fused, config)
response, timing = await async_timed_result("Rerank", apply_rerank(client, query, fused, config))
else:
response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
timings.append(timing)
@@ -172,11 +169,15 @@ def skip_phrase_matches() -> list[HydratedPhraseMatch]:
return []
def query_phrase_matches(engine: Engine, query: str, config: EbookSearchConfig) -> list[HydratedPhraseMatch]:
async def query_phrase_matches(
engine: AsyncEngine,
query: str,
config: EbookSearchConfig,
) -> list[HydratedPhraseMatch]:
"""Detect protected phrases in a query without making search fail when phrase tables are unavailable."""
try:
with Session(engine) as session:
return detect_protected_phrases_for_query(session, query, config)
async with AsyncSession(engine) as session:
return await detect_protected_phrases_for_query(session, query, config)
except SQLAlchemyError as error:
logger.warning("ebook_protected_phrase_detection_unavailable error=%s", error)
return []
@@ -188,8 +189,8 @@ def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchRes
return candidates
def apply_phrase_mention_boosts(
engine: Engine,
async def apply_phrase_mention_boosts(
engine: AsyncEngine,
candidates: list[SearchResult],
phrase_matches: Sequence[HydratedPhraseMatch],
phrase_hit_boost: float,
@@ -201,8 +202,8 @@ def apply_phrase_mention_boosts(
chunk_ids = [candidate.chunk_id for candidate in candidates]
try:
with Session(engine) as session:
phrase_hits = phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
async with AsyncSession(engine) as session:
phrase_hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
except SQLAlchemyError as error:
logger.warning("ebook_phrase_boost_unavailable error=%s", error)
return candidates
@@ -241,30 +242,21 @@ def phrase_rank_source(rank_source: str, phrase_hit_count: int) -> str:
return f"{rank_source} + phrases"
def parallel_retrieval(
engine: Engine,
async def parallel_retrieval(
engine: AsyncEngine,
client: httpx.AsyncClient,
query: str,
config: EbookSearchConfig,
) -> RetrievalResponse:
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions."""
with ThreadPoolExecutor(max_workers=2, thread_name_prefix="ebook-search") as executor:
vector_future = executor.submit(
timed_result,
"Embedding + vector search",
vector_candidates,
engine,
query,
config,
)
bm25_future = executor.submit(
timed_result,
"BM25 search",
bm25_candidates,
query,
config,
)
vector_results, vector_timing = vector_future.result()
lexical_results, lexical_timing = bm25_future.result()
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions.
BM25 scoring is pure CPU work over the cached corpus, so it runs in a worker thread
instead of on the event loop.
"""
(vector_results, vector_timing), (lexical_results, lexical_timing) = await asyncio.gather(
async_timed_result("Embedding + vector search", vector_candidates(engine, client, query, config)),
async_timed_result("BM25 search", asyncio.to_thread(bm25_candidates, query, config)),
)
logger.info(
"ebook_parallel_retrieval_complete vector_candidates=%s lexical_candidates=%s",
@@ -291,13 +283,14 @@ def skip_rerank(
return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
def apply_rerank(
async def apply_rerank(
client: httpx.AsyncClient,
query: str,
candidates: list[SearchResult],
config: EbookSearchConfig,
) -> SearchResponse:
"""Rerank already-fused hybrid candidates."""
reranked = rerank_chunks(query, candidates[: config.rerank.candidates], config.rerank)
reranked = await rerank_chunks(client, query, candidates[: config.rerank.candidates], config.rerank)
logger.info(
"ebook_rerank_complete input_candidates=%s returned=%s",
min(len(candidates), config.rerank.candidates),
@@ -310,10 +303,17 @@ def apply_rerank(
)
def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) -> list[SearchResult]:
async def vector_candidates(
engine: AsyncEngine,
client: httpx.AsyncClient,
query: str,
config: EbookSearchConfig,
) -> list[SearchResult]:
"""Return pgvector cosine candidates for a natural-language query."""
with Session(engine) as session:
model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
async with AsyncSession(engine) as session:
model = await session.scalar(
select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model)
)
if model is None:
msg = f"Embedding model is not registered: {config.embedding_model}"
raise ValueError(msg)
@@ -323,7 +323,7 @@ def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) ->
msg = f"Model row dimension {model.dimension} does not match configured dimension {expected_dimension}"
raise ValueError(msg)
embedding = embed_query(query, config)
embedding = await embed_query(client, query, config)
limit = max(config.rerank.candidates, config.top_k) * config.vector_candidate_multiplier
embedding_table = get_embedding_table(model.dimension)
@@ -349,7 +349,7 @@ def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) ->
.order_by(distance)
.limit(limit)
)
rows = session.execute(statement).mappings()
rows = (await session.execute(statement)).mappings()
results = [search_result_from_row(row) for row in rows]
logger.info(
"ebook_vector_search_complete model=%s dimension=%s candidates=%s",