Richie
c7cd63f8e4
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
74 lines
2.5 KiB
Python
74 lines
2.5 KiB
Python
"""Background BM25 refresh tasks for the web app.
|
|
|
|
The refresh is scheduled on the event loop instead of a thread because the async psycopg
|
|
driver only works from the loop; a bare thread cannot open a session on the async engine.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from typing import TYPE_CHECKING
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from python.ebook_search.bm25_corpus import load_bm25_corpus, refresh_bm25_corpus
|
|
|
|
if TYPE_CHECKING:
|
|
from fastapi import FastAPI
|
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
|
from python.ebook_search.config import EbookSearchConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def schedule_bm25_refresh(app: FastAPI) -> None:
|
|
"""Schedule a delayed BM25 corpus refresh, replacing any pending refresh.
|
|
|
|
Only called from route handlers, so a running event loop is guaranteed.
|
|
"""
|
|
cancel_bm25_refresh(app)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
|
|
def start_refresh() -> None:
|
|
app.state.bm25_refresh_task = loop.create_task(refresh_bm25_for_app(app))
|
|
|
|
app.state.bm25_refresh_timer = loop.call_later(app.state.config.bm25_refresh_delay_seconds, start_refresh)
|
|
logger.info(
|
|
"ebook_bm25_refresh_scheduled delay_seconds=%s",
|
|
app.state.config.bm25_refresh_delay_seconds,
|
|
)
|
|
|
|
|
|
def cancel_bm25_refresh(app: FastAPI) -> None:
|
|
"""Cancel any pending BM25 corpus refresh timer and in-flight refresh task."""
|
|
existing_timer = getattr(app.state, "bm25_refresh_timer", None)
|
|
if existing_timer is not None:
|
|
existing_timer.cancel()
|
|
app.state.bm25_refresh_timer = None
|
|
logger.info("ebook_bm25_refresh_cancelled")
|
|
|
|
existing_task = getattr(app.state, "bm25_refresh_task", None)
|
|
if existing_task is not None:
|
|
if not existing_task.done():
|
|
existing_task.cancel()
|
|
app.state.bm25_refresh_task = None
|
|
|
|
|
|
async def refresh_bm25_for_app(app: FastAPI) -> None:
|
|
"""Refresh the BM25 corpus using the app engine and config."""
|
|
try:
|
|
await refresh_bm25_for_engine(app.state.engine, app.state.config)
|
|
except Exception:
|
|
logger.exception("ebook_bm25_refresh_failed")
|
|
|
|
|
|
async def refresh_bm25_for_engine(engine: AsyncEngine, config: EbookSearchConfig) -> None:
|
|
"""Refresh the BM25 corpus using an async SQLAlchemy engine."""
|
|
async with AsyncSession(engine) as session:
|
|
await refresh_bm25_corpus(session, config)
|
|
load_bm25_corpus.cache_clear()
|
|
logger.info("ebook_bm25_corpus_cache_cleared_after_refresh")
|