Richie
2706c4417d
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
99 lines
3.4 KiB
Python
99 lines
3.4 KiB
Python
"""FastAPI HTMX app for EPUB search."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from contextlib import asynccontextmanager
|
|
from typing import TYPE_CHECKING, Annotated
|
|
|
|
import httpx
|
|
import typer
|
|
import uvicorn
|
|
from fastapi import FastAPI
|
|
from fastapi.staticfiles import StaticFiles
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from python.common import configure_logger
|
|
from python.ebook_search.api.bm25_tasks import cancel_bm25_refresh
|
|
from python.ebook_search.api.routes import admin_router, health_router, page_router, search_router
|
|
from python.ebook_search.api.web import STATIC_DIR
|
|
from python.ebook_search.bm25_corpus import ensure_bm25_corpus
|
|
from python.ebook_search.config import load_config
|
|
from python.ebook_search.protected_phrases.pool import shutdown_extraction_pool
|
|
from python.fastapi_tools import ZstdMiddleware
|
|
from python.orm.common import get_async_postgres_engine
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
|
|
"""Manage application startup and shutdown resources."""
|
|
logger.info("ebook_search_startup")
|
|
config = load_config()
|
|
app.state.config = config
|
|
logger.info(
|
|
"ebook_search_config_loaded top_k=%s embedding_model=%s embedding_base_url=%s vllm_base_url=%s "
|
|
"rerank_enabled=%s phrase_matching_enabled=%s answer_enabled=%s library_paths=%s",
|
|
config.top_k,
|
|
config.embedding_model,
|
|
config.embedding_base_url,
|
|
config.vllm_base_url,
|
|
config.rerank.enabled,
|
|
config.phrase_matching_enabled,
|
|
config.answer_enabled,
|
|
len(config.library_paths),
|
|
)
|
|
if not config.library_paths:
|
|
logger.warning("ebook_search_no_library_paths_configured")
|
|
# Concurrent phrase judging opens one session per book worker on this engine, so size the pool
|
|
# to cover those plus headroom for ordinary web requests.
|
|
app.state.engine = get_async_postgres_engine(
|
|
name="RICHIE",
|
|
vector_engine=True,
|
|
pool_size=config.phrase_judge_book_workers + 10,
|
|
)
|
|
app.state.http_client = httpx.AsyncClient()
|
|
async with AsyncSession(app.state.engine, expire_on_commit=False) as session:
|
|
await ensure_bm25_corpus(session, config)
|
|
try:
|
|
yield
|
|
finally:
|
|
logger.info("ebook_search_shutdown")
|
|
cancel_bm25_refresh(app)
|
|
shutdown_extraction_pool()
|
|
await app.state.http_client.aclose()
|
|
await app.state.engine.dispose()
|
|
|
|
|
|
def create_app() -> FastAPI:
|
|
"""Create the EPUB search web app."""
|
|
app = FastAPI(title="EPUB Search", lifespan=lifespan)
|
|
app.add_middleware(ZstdMiddleware)
|
|
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
|
|
|
|
app.include_router(admin_router)
|
|
app.include_router(health_router)
|
|
app.include_router(page_router)
|
|
app.include_router(search_router)
|
|
|
|
return app
|
|
|
|
|
|
def serve(
|
|
host: Annotated[str, typer.Option("--host", "-h", help="Host to bind to")] = "127.0.0.1",
|
|
port: Annotated[int, typer.Option("--port", "-p", help="Port to bind to")] = 8070,
|
|
log_level: Annotated[str, typer.Option("--log-level", "-l", help="Log level")] = "INFO",
|
|
) -> None:
|
|
"""Start the EPUB search server."""
|
|
configure_logger(log_level)
|
|
uvicorn.run(create_app(), host=host, port=port)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
typer.run(serve)
|