diff --git a/python/ebook_search/answer.py b/python/ebook_search/answer.py index 4b85b21..ba7f6aa 100644 --- a/python/ebook_search/answer.py +++ b/python/ebook_search/answer.py @@ -8,13 +8,20 @@ from typing import TYPE_CHECKING from python.ebook_search.llm_interface import request_chat_completion if TYPE_CHECKING: + import httpx + from python.ebook_search.config import EbookSearchConfig from python.ebook_search.search import SearchResult logger = logging.getLogger(__name__) -def answer_query(query: str, results: list[SearchResult], config: EbookSearchConfig) -> str: +async def answer_query( + client: httpx.AsyncClient, + query: str, + results: list[SearchResult], + config: EbookSearchConfig, +) -> str: """Answer a question using only retrieved chunks.""" if not config.answer_enabled: logger.info("ebook_answer_skipped_disabled") @@ -35,7 +42,8 @@ def answer_query(query: str, results: list[SearchResult], config: EbookSearchCon f"[{index}] {result.source_title}{' - ' + result.chapter_title if result.chapter_title else ''}\n{result.text}" for index, result in enumerate(results, start=1) ) - content = request_chat_completion( + content = await request_chat_completion( + client, config, [ { diff --git a/python/ebook_search/api/bm25_tasks.py b/python/ebook_search/api/bm25_tasks.py index ff24b85..3477b6c 100644 --- a/python/ebook_search/api/bm25_tasks.py +++ b/python/ebook_search/api/bm25_tasks.py @@ -1,18 +1,22 @@ -"""Background BM25 refresh tasks for the web app.""" +"""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 threading import Timer from typing import TYPE_CHECKING -from sqlalchemy.orm import Session +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.engine import Engine + from sqlalchemy.ext.asyncio import AsyncEngine from python.ebook_search.config import EbookSearchConfig @@ -20,15 +24,18 @@ logger = logging.getLogger(__name__) def schedule_bm25_refresh(app: FastAPI) -> None: - """Schedule a delayed BM25 corpus refresh, replacing any pending refresh.""" - existing_timer = getattr(app.state, "bm25_refresh_timer", None) - if existing_timer is not None: - existing_timer.cancel() + """Schedule a delayed BM25 corpus refresh, replacing any pending refresh. - timer = Timer(app.state.config.bm25_refresh_delay_seconds, refresh_bm25_for_app, args=(app,)) - timer.daemon = True - timer.start() - app.state.bm25_refresh_timer = timer + 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, @@ -36,25 +43,31 @@ def schedule_bm25_refresh(app: FastAPI) -> None: def cancel_bm25_refresh(app: FastAPI) -> None: - """Cancel any pending BM25 corpus refresh.""" + """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 -def refresh_bm25_for_app(app: FastAPI) -> None: + +async def refresh_bm25_for_app(app: FastAPI) -> None: """Refresh the BM25 corpus using the app engine and config.""" try: - refresh_bm25_for_engine(app.state.engine, app.state.config) + await refresh_bm25_for_engine(app.state.engine, app.state.config) except Exception: logger.exception("ebook_bm25_refresh_failed") -def refresh_bm25_for_engine(engine: Engine, config: EbookSearchConfig) -> None: - """Refresh the BM25 corpus using a SQLAlchemy engine.""" - with Session(engine) as session: - refresh_bm25_corpus(session, config) +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") diff --git a/python/ebook_search/api/dependencies.py b/python/ebook_search/api/dependencies.py index e5a89fa..508dfd7 100644 --- a/python/ebook_search/api/dependencies.py +++ b/python/ebook_search/api/dependencies.py @@ -4,8 +4,9 @@ from __future__ import annotations from typing import Annotated +import httpx from fastapi import Depends, Request -from sqlalchemy.engine import Engine +from sqlalchemy.ext.asyncio import AsyncEngine from python.ebook_search.config import EbookSearchConfig @@ -15,10 +16,16 @@ def get_config(request: Request) -> EbookSearchConfig: return request.app.state.config -def get_engine(request: Request) -> Engine: +def get_engine(request: Request) -> AsyncEngine: """Get the database engine from app state.""" return request.app.state.engine +def get_http_client(request: Request) -> httpx.AsyncClient: + """Get the shared LLM HTTP client from app state.""" + return request.app.state.http_client + + AppConfig = Annotated[EbookSearchConfig, Depends(get_config)] -AppEngine = Annotated[Engine, Depends(get_engine)] +AppEngine = Annotated[AsyncEngine, Depends(get_engine)] +AppHttpClient = Annotated[httpx.AsyncClient, Depends(get_http_client)] diff --git a/python/ebook_search/api/judge_tasks.py b/python/ebook_search/api/judge_tasks.py new file mode 100644 index 0000000..7bc07b0 --- /dev/null +++ b/python/ebook_search/api/judge_tasks.py @@ -0,0 +1,131 @@ +"""Background phrase-judging tasks for the web app. + +Judging a book sends one LLM request per candidate phrase, which can take minutes, so it must +not run inside the request where it would block the UI. Judgments run as async FastAPI +background tasks, awaited on the event loop after the response is sent, and are tracked per +book in app state so a second judge request for a book that is already being judged is +rejected instead of doubling the work. + +State is loop-confined: every read and mutation happens on the event loop (async route +handlers and async background tasks) and no critical section contains an ``await``, so each +mutation is atomic per loop iteration and no locking is needed. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +from python.ebook_search.protected_phrases.judge_ngrams import judge_candidate_phrases_for_books + +if TYPE_CHECKING: + from fastapi import BackgroundTasks, FastAPI + +logger = logging.getLogger(__name__) + + +@dataclass +class JudgeTaskState: + """Running book judgments and last outcome messages, keyed by book id.""" + + running_book_ids: set[int] = field(default_factory=set) + outcome_messages: dict[int, str] = field(default_factory=dict) + + +def get_judge_task_state(app: FastAPI) -> JudgeTaskState: + """Return the app's judge task state, creating it on first use. + + Args: + app (FastAPI): App whose state holds the judge task registry. + + Returns: + JudgeTaskState: The shared judge task state for this app. + """ + state = getattr(app.state, "judge_tasks", None) + if state is None: + state = JudgeTaskState() + app.state.judge_tasks = state + return state + + +def start_book_phrase_judgment(app: FastAPI, background_tasks: BackgroundTasks, source_id: int) -> bool: + """Queue judging of one book's candidate phrases as a FastAPI background task. + + The book is claimed before the response returns, so a repeated judge request cannot queue + a second run while one is pending or running. + + Args: + app (FastAPI): App supplying the engine, config, and judge task state. + background_tasks (BackgroundTasks): Request's background tasks to queue the judgment on. + source_id (int): Book to judge candidates for. + + Returns: + bool: True when a judgment was queued, False when one is already running for this book. + """ + state = get_judge_task_state(app) + if source_id in state.running_book_ids: + logger.info("ebook_book_phrase_judgment_already_running source_id=%s", source_id) + return False + state.running_book_ids.add(source_id) + state.outcome_messages.pop(source_id, None) + background_tasks.add_task(judge_book_phrases_for_app, app, source_id) + logger.info("ebook_book_phrase_judgment_queued source_id=%s", source_id) + return True + + +async def judge_book_phrases_for_app(app: FastAPI, source_id: int) -> None: + """Judge one book using the app engine and config, recording the outcome message. + + Args: + app (FastAPI): App supplying the engine, config, and judge task state. + source_id (int): Book to judge candidates for. + """ + state = get_judge_task_state(app) + try: + result = await judge_candidate_phrases_for_books(app.state.engine, app.state.config, source_ids=[source_id]) + logger.info( + "ebook_book_phrase_judgment_complete source_id=%s judged=%s protected=%s mentions=%s failed=%s", + source_id, + result.candidates_judged, + result.protected_phrases, + result.phrase_mentions, + result.books_failed, + ) + if result.books_failed: + message = "Judging failed; see server logs for details" + else: + message = ( + f"Judged {result.candidates_judged} candidates; {result.protected_phrases} protected phrases promoted" + ) + except Exception: + logger.exception("ebook_book_phrase_judgment_task_failed source_id=%s", source_id) + message = "Judging failed; see server logs for details" + state.running_book_ids.discard(source_id) + state.outcome_messages[source_id] = message + + +def is_judging_book(app: FastAPI, source_id: int) -> bool: + """Report whether a judgment is currently queued or running for one book. + + Args: + app (FastAPI): App supplying the judge task state. + source_id (int): Book to check. + + Returns: + bool: True while the book's judgment is pending or running. + """ + return source_id in get_judge_task_state(app).running_book_ids + + +def pop_book_judgment_outcome(app: FastAPI, source_id: int) -> str | None: + """Return and clear the outcome message from one book's last finished judgment. + + Args: + app (FastAPI): App supplying the judge task state. + source_id (int): Book to fetch the outcome for. + + Returns: + str | None: The outcome message, or None when there is nothing new to report. + """ + return get_judge_task_state(app).outcome_messages.pop(source_id, None) diff --git a/python/ebook_search/api/main.py b/python/ebook_search/api/main.py index 87e3316..4b286d1 100644 --- a/python/ebook_search/api/main.py +++ b/python/ebook_search/api/main.py @@ -6,11 +6,12 @@ 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.orm import Session +from sqlalchemy.ext.asyncio import AsyncSession from python.common import configure_logger from python.ebook_search.api.bm25_tasks import cancel_bm25_refresh @@ -18,8 +19,9 @@ from python.ebook_search.api.routes import admin_router, health_router, page_rou 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_postgres_engine +from python.orm.common import get_async_postgres_engine if TYPE_CHECKING: from collections.abc import AsyncIterator @@ -48,15 +50,24 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: ) if not config.library_paths: logger.warning("ebook_search_no_library_paths_configured") - app.state.engine = get_postgres_engine(name="RICHIE", vector_engine=True) - with Session(app.state.engine) as session: - ensure_bm25_corpus(session, config) + # 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) - app.state.engine.dispose() + shutdown_extraction_pool() + await app.state.http_client.aclose() + await app.state.engine.dispose() def create_app() -> FastAPI: diff --git a/python/ebook_search/api/routes/admin.py b/python/ebook_search/api/routes/admin.py index 737ec74..6860609 100644 --- a/python/ebook_search/api/routes/admin.py +++ b/python/ebook_search/api/routes/admin.py @@ -8,15 +8,18 @@ from fastapi import APIRouter, Request from fastapi.responses import HTMLResponse from python.ebook_search.api.bm25_tasks import schedule_bm25_refresh -from python.ebook_search.api.dependencies import ( - AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime +from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime + AppConfig, + AppEngine, + AppHttpClient, ) from python.ebook_search.api.web import templates from python.ebook_search.embeddings import embed_missing_chunks, embedding_model_stats from python.ebook_search.ingest import ingest_configured_paths from python.ebook_search.protected_phrases.generate_ngrams import generate_candidate_phrases_for_books from python.ebook_search.protected_phrases.judge_ngrams import judge_candidate_phrases_for_books -from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime +from python.ebook_search.protected_phrases.store import book_ids_pending_first_judgment, corpus_phrase_stats +from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime logger = logging.getLogger(__name__) @@ -24,19 +27,29 @@ router = APIRouter(prefix="/admin") @router.get("", response_class=HTMLResponse) -def admin(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse: +async def admin(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse: """Render the admin page.""" - stats = embedding_model_stats(session) - logger.info("ebook_admin_page_loaded models=%s", len(stats)) - return templates.TemplateResponse(request, "admin.html", {"config": config, "stats": stats}) + stats = await embedding_model_stats(session) + phrase_stats = await corpus_phrase_stats(session) + logger.info( + "ebook_admin_page_loaded models=%s candidate_phrases=%s protected_phrases=%s", + len(stats), + phrase_stats.candidate_phrases, + phrase_stats.protected_phrases, + ) + return templates.TemplateResponse( + request, + "admin.html", + {"config": config, "stats": stats, "phrase_stats": phrase_stats}, + ) @router.post("/scan", response_class=HTMLResponse) -def scan_library(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse: +async def scan_library(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse: """Scan configured library paths for EPUB changes.""" try: - count = ingest_configured_paths(session, config) - session.commit() + count = await ingest_configured_paths(session, config) + await session.commit() except Exception as error: logger.exception("ebook_admin_scan_failed") return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) @@ -47,48 +60,119 @@ def scan_library(request: Request, config: AppConfig, session: DbSession) -> HTM return templates.TemplateResponse(request, "partials/admin_status.html", {"message": f"Indexed {count} EPUBs"}) -@router.post("/generate-ngrams", response_class=HTMLResponse) -def generate_ngrams(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse: - """Generate candidate n-grams for indexed books without LLM judging.""" +@router.post("/phrases/generate-all", response_class=HTMLResponse) +async def generate_all_phrases(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse: + """Regenerate candidate phrases for every indexed book without LLM judging.""" + return await run_phrase_generation(request, config, session, only_missing=False) + + +@router.post("/phrases/generate-missing", response_class=HTMLResponse) +async def generate_missing_phrases(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse: + """Generate candidate phrases only for books that have none yet.""" + return await run_phrase_generation(request, config, session, only_missing=True) + + +async def run_phrase_generation( + request: Request, + config: AppConfig, + session: AsyncDbSession, + *, + only_missing: bool, +) -> HTMLResponse: + """Run candidate phrase generation and render the outcome as an admin status partial. + + Args: + request (Request): Current request, for template rendering. + config (AppConfig): Runtime phrase-tuning settings. + session (AsyncDbSession): Active database session. + only_missing (bool): Only generate for books without candidates instead of every book. + + Returns: + HTMLResponse: Status partial describing the generation outcome. + """ try: - result = generate_candidate_phrases_for_books(session, config) - session.commit() + result = await generate_candidate_phrases_for_books(session, config, only_missing=only_missing) + await session.commit() except Exception as error: - session.rollback() - logger.exception("ebook_admin_generate_ngrams_failed") + await session.rollback() + logger.exception("ebook_admin_generate_phrases_failed only_missing=%s", only_missing) return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) logger.info( - "ebook_admin_generate_ngrams_complete books_seen=%s books_built=%s candidates=%s", + "ebook_admin_generate_phrases_complete only_missing=%s books_seen=%s books_built=%s candidates=%s", + only_missing, result.books_seen, result.books_built, result.candidate_phrases, ) + if only_missing and result.books_seen == 0: + return templates.TemplateResponse( + request, + "partials/admin_status.html", + {"message": "All books already have candidate phrases"}, + ) return templates.TemplateResponse( request, "partials/admin_status.html", { "message": ( - f"Generated n-grams for {result.books_built} of {result.books_seen} books; " + f"Generated phrases for {result.books_built} of {result.books_seen} books; " f"{result.candidate_phrases} candidates stored" ) }, ) -@router.post("/judge-ngrams", response_class=HTMLResponse) -def judge_ngrams(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse: - """Judge stored candidate n-grams and promote accepted protected phrases.""" +@router.post("/phrases/judge-all", response_class=HTMLResponse) +async def judge_all_phrases(request: Request, engine: AppEngine, config: AppConfig) -> HTMLResponse: + """Judge unjudged candidate phrases across every indexed book.""" + return await run_phrase_judgment(request, engine, config, source_ids=None) + + +@router.post("/phrases/judge-missing", response_class=HTMLResponse) +async def judge_missing_phrases( + request: Request, + engine: AppEngine, + config: AppConfig, + session: AsyncDbSession, +) -> HTMLResponse: + """Judge candidate phrases only for books where judging has never run.""" + source_ids = await book_ids_pending_first_judgment(session) + if not source_ids: + return templates.TemplateResponse( + request, + "partials/admin_status.html", + {"message": "All books with candidate phrases have been judged"}, + ) + return await run_phrase_judgment(request, engine, config, source_ids=source_ids) + + +async def run_phrase_judgment( + request: Request, + engine: AppEngine, + config: AppConfig, + *, + source_ids: list[int] | None, +) -> HTMLResponse: + """Run LLM judging for candidate phrases and render the outcome as an admin status partial. + + Args: + request (Request): Current request, for template rendering. + engine (AppEngine): Engine used to open per-book judging sessions. + config (AppConfig): Runtime phrase-tuning settings. + source_ids (list[int] | None): Books to judge; ``None`` judges every indexed book. + + Returns: + HTMLResponse: Status partial describing the judging outcome. + """ try: - result = judge_candidate_phrases_for_books(session, config) - session.commit() + result = await judge_candidate_phrases_for_books(engine, config, source_ids=source_ids) except Exception as error: - session.rollback() - logger.exception("ebook_admin_judge_ngrams_failed") + logger.exception("ebook_admin_judge_phrases_failed") return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) logger.info( - "ebook_admin_judge_ngrams_complete books_seen=%s books_judged=%s books_failed=%s candidates_judged=%s " + "ebook_admin_judge_phrases_complete books_seen=%s books_judged=%s books_failed=%s candidates_judged=%s " "protected=%s mentions=%s", result.books_seen, result.books_judged, @@ -112,11 +196,16 @@ def judge_ngrams(request: Request, config: AppConfig, session: DbSession) -> HTM @router.post("/embed-missing", response_class=HTMLResponse) -def embed_missing(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse: +async def embed_missing( + request: Request, + config: AppConfig, + session: AsyncDbSession, + client: AppHttpClient, +) -> HTMLResponse: """Embed chunks missing vectors for the configured model.""" try: - count = embed_missing_chunks(session, config) - session.commit() + count = await embed_missing_chunks(session, client, config) + await session.commit() except Exception as error: logger.exception("ebook_admin_embed_missing_failed") return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) @@ -130,16 +219,21 @@ def embed_missing(request: Request, config: AppConfig, session: DbSession) -> HT @router.post("/embed-all", response_class=HTMLResponse) -def embed_all(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse: +async def embed_all( + request: Request, + config: AppConfig, + session: AsyncDbSession, + client: AppHttpClient, +) -> HTMLResponse: """Embed all chunks missing vectors in fixed-size batches.""" total = 0 batches = 0 try: while True: - count = embed_missing_chunks(session, config) + count = await embed_missing_chunks(session, client, config) if count == 0: break - session.commit() + await session.commit() total += count batches += 1 logger.info( diff --git a/python/ebook_search/api/routes/health.py b/python/ebook_search/api/routes/health.py index fb7c097..f5726fd 100644 --- a/python/ebook_search/api/routes/health.py +++ b/python/ebook_search/api/routes/health.py @@ -11,15 +11,17 @@ from fastapi.responses import JSONResponse from sqlalchemy import literal, select from sqlalchemy.exc import SQLAlchemyError -from python.ebook_search.api.dependencies import ( - AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime +from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime + AppConfig, + AppHttpClient, ) from python.ebook_search.bm25_corpus import bm25_index_exists, bm25_index_path, read_bm25_manifest from python.ebook_search.llm_interface import check_chat_endpoint, check_embedding_endpoint -from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime +from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime if TYPE_CHECKING: - from sqlalchemy.orm import Session + import httpx + from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.config import EbookSearchConfig @@ -29,17 +31,17 @@ router = APIRouter() @router.get("/health") -def health() -> dict[str, str]: +async def health() -> dict[str, str]: """Liveness probe that returns ok without touching dependencies.""" return {"status": "ok"} @router.get("/ready") -def ready(config: AppConfig, session: DbSession) -> JSONResponse: +async def ready(config: AppConfig, session: AsyncDbSession, client: AppHttpClient) -> JSONResponse: """Readiness probe reporting database, embedding endpoint, and BM25 index status.""" - database_ok = check_database(session) - embedding_ok = check_embedding_endpoint(config) - chat_status = chat_endpoint_status(config) + database_ok = await check_database(session) + embedding_ok = await check_embedding_endpoint(client, config) + chat_status = await chat_endpoint_status(client, config) bm25_status = check_bm25_status(config) checks = { @@ -69,17 +71,17 @@ def ready(config: AppConfig, session: DbSession) -> JSONResponse: return JSONResponse(content={"status": status, "checks": checks}, status_code=status_code) -def chat_endpoint_status(config: EbookSearchConfig) -> str: +async def chat_endpoint_status(client: httpx.AsyncClient, config: EbookSearchConfig) -> str: """Return the answering chat endpoint status, or disabled when answers are off.""" if not config.answer_enabled: return "disabled" - return "ok" if check_chat_endpoint(config) else "fail" + return "ok" if await check_chat_endpoint(client, config) else "fail" -def check_database(session: Session) -> bool: +async def check_database(session: AsyncSession) -> bool: """Return whether the database answers a trivial query.""" try: - session.execute(select(literal(1))) + await session.execute(select(literal(1))) except SQLAlchemyError as error: logger.warning("ebook_ready_database_unavailable error=%s", error) return False diff --git a/python/ebook_search/api/routes/page.py b/python/ebook_search/api/routes/page.py index 4608f39..47858ee 100644 --- a/python/ebook_search/api/routes/page.py +++ b/python/ebook_search/api/routes/page.py @@ -3,18 +3,23 @@ from __future__ import annotations import logging +from typing import TYPE_CHECKING -from fastapi import APIRouter, HTTPException, Request +from fastapi import APIRouter, BackgroundTasks, HTTPException, Request from fastapi.responses import HTMLResponse, RedirectResponse from sqlalchemy import func, select from python.ebook_search.api.dependencies import ( AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime ) +from python.ebook_search.api.judge_tasks import is_judging_book, pop_book_judgment_outcome, start_book_phrase_judgment from python.ebook_search.api.web import templates from python.ebook_search.protected_phrases.generate_ngrams import recalculate_candidate_phrases_for_book -from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime -from python.orm.richie import EbookCandidatePhrase, EbookProtectedPhrase, EbookSource +from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime +from python.orm.richie import EbookCandidatePhrase, EbookChapter, EbookChunk, EbookProtectedPhrase, EbookSource + +if TYPE_CHECKING: + from sqlalchemy.ext.asyncio import AsyncSession logger = logging.getLogger(__name__) @@ -22,30 +27,41 @@ router = APIRouter() @router.get("/", response_class=HTMLResponse) -def index(request: Request, config: AppConfig) -> HTMLResponse: +async def index(request: Request, config: AppConfig) -> HTMLResponse: """Render the search page.""" return templates.TemplateResponse(request, "search.html", {"config": config}) @router.get("/books", response_class=HTMLResponse) -def books(request: Request, session: DbSession) -> HTMLResponse: +async def books(request: Request, session: AsyncDbSession) -> HTMLResponse: """Render the indexed books page.""" - sources = list(session.scalars(select(EbookSource).order_by(EbookSource.title)).all()) + sources = list((await session.scalars(select(EbookSource).order_by(EbookSource.title))).all()) logger.info("ebook_books_page_loaded count=%s", len(sources)) return templates.TemplateResponse(request, "books.html", {"sources": sources}) -def get_candidate_count(session: DbSession, book_id: int) -> int: +async def get_chapter_count(session: AsyncSession, book_id: int) -> int: + """Return the number of indexed chapters for one book.""" + return await session.scalar(select(func.count(EbookChapter.id)).where(EbookChapter.source_id == book_id)) or 0 + + +async def get_chunk_count(session: AsyncSession, book_id: int) -> int: + """Return the number of indexed chunks for one book.""" + return await session.scalar(select(func.count(EbookChunk.id)).where(EbookChunk.source_id == book_id)) or 0 + + +async def get_candidate_count(session: AsyncSession, book_id: int) -> int: """Return the number of indexed candidates for one book.""" return ( - session.scalar(select(func.count(EbookCandidatePhrase.id)).where(EbookCandidatePhrase.book_id == book_id)) or 0 + await session.scalar(select(func.count(EbookCandidatePhrase.id)).where(EbookCandidatePhrase.book_id == book_id)) + or 0 ) -def get_judged_candidate_count(session: DbSession, book_id: int) -> int: +async def get_judged_candidate_count(session: AsyncSession, book_id: int) -> int: """Return the number of judged candidates for one book.""" return ( - session.scalar( + await session.scalar( select(func.count(EbookCandidatePhrase.id)).where( EbookCandidatePhrase.book_id == book_id, EbookCandidatePhrase.llm_judged.is_(True), @@ -55,17 +71,18 @@ def get_judged_candidate_count(session: DbSession, book_id: int) -> int: ) -def get_protected_count(session: DbSession, book_id: int) -> int: +async def get_protected_count(session: AsyncSession, book_id: int) -> int: """Return the number of protected phrases for one book.""" return ( - session.scalar(select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)) or 0 + await session.scalar(select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)) + or 0 ) -def get_candidates(session: DbSession, book_id: int) -> list[EbookCandidatePhrase]: +async def get_candidates(session: AsyncSession, book_id: int) -> list[EbookCandidatePhrase]: """Return the indexed candidates for one book.""" return list( - session.scalars( + await session.scalars( select(EbookCandidatePhrase) .where(EbookCandidatePhrase.book_id == book_id) .order_by(EbookCandidatePhrase.candidate_score.desc()) @@ -74,10 +91,10 @@ def get_candidates(session: DbSession, book_id: int) -> list[EbookCandidatePhras ) -def get_protected_phrases(session: DbSession, book_id: int) -> list[EbookProtectedPhrase]: +async def get_protected_phrases(session: AsyncSession, book_id: int) -> list[EbookProtectedPhrase]: """Return the protected phrases for one book.""" return list( - session.scalars( + await session.scalars( select(EbookProtectedPhrase) .where(EbookProtectedPhrase.book_id == book_id) .order_by(EbookProtectedPhrase.importance.desc()) @@ -87,21 +104,27 @@ def get_protected_phrases(session: DbSession, book_id: int) -> list[EbookProtect @router.get("/books/{source_id}", response_class=HTMLResponse) -def book_detail(source_id: int, request: Request, session: DbSession) -> HTMLResponse: +async def book_detail(source_id: int, request: Request, session: AsyncDbSession) -> HTMLResponse: """Render details for one indexed book.""" - source = session.get(EbookSource, source_id) + source = await session.get(EbookSource, source_id) phrase_status_message = None recalculated = request.query_params.get("phrases_recalculated") if recalculated is not None: phrase_status_message = f"Recalculated phrases; {recalculated} candidates generated" + judgment_outcome = pop_book_judgment_outcome(request.app, source_id) + if judgment_outcome is not None: + phrase_status_message = judgment_outcome + judging_in_progress = is_judging_book(request.app, source_id) + if judging_in_progress: + phrase_status_message = "Judging candidate phrases in the background; refresh to see progress" if source is not None: - chapter_count = len(source.chapters) - chunk_count = len(source.chunks) - candidate_count = get_candidate_count(session, source.id) - judged_candidate_count = get_judged_candidate_count(session, source.id) - protected_count = get_protected_count(session, source.id) - candidates = get_candidates(session, source.id) - protected_phrases = get_protected_phrases(session, source.id) + chapter_count = await get_chapter_count(session, source.id) + chunk_count = await get_chunk_count(session, source.id) + candidate_count = await get_candidate_count(session, source.id) + judged_candidate_count = await get_judged_candidate_count(session, source.id) + protected_count = await get_protected_count(session, source.id) + candidates = await get_candidates(session, source.id) + protected_phrases = await get_protected_phrases(session, source.id) else: chapter_count = 0 chunk_count = 0 @@ -129,6 +152,7 @@ def book_detail(source_id: int, request: Request, session: DbSession) -> HTMLRes "chapter_count": chapter_count, "chunk_count": chunk_count, "judged_candidate_count": judged_candidate_count, + "judging_in_progress": judging_in_progress, "protected_count": protected_count, "protected_phrases": protected_phrases, "phrase_status_message": phrase_status_message, @@ -138,13 +162,13 @@ def book_detail(source_id: int, request: Request, session: DbSession) -> HTMLRes @router.post("/books/{source_id}/recalculate-phrases") -def recalculate_book_phrases(source_id: int, config: AppConfig, session: DbSession) -> RedirectResponse: +async def recalculate_book_phrases(source_id: int, config: AppConfig, session: AsyncDbSession) -> RedirectResponse: """Clear and regenerate candidate phrases for one indexed book.""" - source = session.get(EbookSource, source_id) + source = await session.get(EbookSource, source_id) if source is None: raise HTTPException(status_code=404, detail="Book not found") - result = recalculate_candidate_phrases_for_book(session, source, config) + result = await recalculate_candidate_phrases_for_book(session, source, config, use_process_pool=True) logger.info( "ebook_book_phrase_recalculation_complete source_id=%s candidates=%s deleted_candidates=%s " "deleted_protected=%s deleted_aliases=%s deleted_mentions=%s", @@ -159,3 +183,20 @@ def recalculate_book_phrases(source_id: int, config: AppConfig, session: DbSessi url=f"/books/{source_id}?phrases_recalculated={result.candidate_phrases}", status_code=303, ) + + +@router.post("/books/{source_id}/judge-phrases") +async def judge_book_phrases( + source_id: int, + request: Request, + background_tasks: BackgroundTasks, + session: AsyncDbSession, +) -> RedirectResponse: + """Queue background judging of one book's candidate phrases and return immediately.""" + source = await session.get(EbookSource, source_id) + if source is None: + raise HTTPException(status_code=404, detail="Book not found") + + started = start_book_phrase_judgment(request.app, background_tasks, source.id) + logger.info("ebook_book_phrase_judgment_requested source_id=%s started=%s", source_id, started) + return RedirectResponse(url=f"/books/{source_id}", status_code=303) diff --git a/python/ebook_search/api/routes/search.py b/python/ebook_search/api/routes/search.py index 6c7f8a6..0cc95b4 100644 --- a/python/ebook_search/api/routes/search.py +++ b/python/ebook_search/api/routes/search.py @@ -14,6 +14,7 @@ from python.ebook_search.answer import answer_query from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolves these annotated dependencies at runtime AppConfig, AppEngine, + AppHttpClient, ) from python.ebook_search.api.web import templates from python.ebook_search.guardrails import ( @@ -26,6 +27,8 @@ from python.ebook_search.search import SearchResponse, search_ebooks from python.ebook_search.timing import runtime_step_from_start if TYPE_CHECKING: + import httpx + from python.ebook_search.config import EbookSearchConfig logger = logging.getLogger(__name__) @@ -33,7 +36,8 @@ logger = logging.getLogger(__name__) router = APIRouter() -def build_answer( +async def build_answer( + client: httpx.AsyncClient, query: str, response: SearchResponse, config: EbookSearchConfig, @@ -56,7 +60,7 @@ def build_answer( return answer, True, None try: - answer = answer_query(query, response.results, config) + answer = await answer_query(client, query, response.results, config) except RuntimeError as error: logger.warning("ebook_answer_request_failed_falling_back error=%s", error) return "Answer generation failed. Source chunks are still shown below.", False, None @@ -74,18 +78,20 @@ def build_answer( @router.post("/search", response_class=HTMLResponse) -def search( +async def search( request: Request, config: AppConfig, engine: AppEngine, + client: AppHttpClient, query: Annotated[str, Form()], rerank: Annotated[str | None, Form()] = None, phrase_matching: Annotated[str | None, Form()] = None, ) -> HTMLResponse: """Run a search and render HTMX results.""" try: - response = search_ebooks( + response = await search_ebooks( engine, + client, query, config, rerank=rerank == "true", @@ -96,7 +102,7 @@ def search( return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) answer_start = perf_counter() - answer, low_confidence, citation_report = build_answer(query, response, config) + answer, low_confidence, citation_report = await build_answer(client, query, response, config) answer_step_name = "Answer generation" if config.answer_enabled else "Answer skipped" response = replace( response, diff --git a/python/ebook_search/api/templates/admin.html b/python/ebook_search/api/templates/admin.html index cd2dacc..512da3a 100644 --- a/python/ebook_search/api/templates/admin.html +++ b/python/ebook_search/api/templates/admin.html @@ -8,37 +8,25 @@ head %}
-
- -
-
- -
-
- -
-
- -

Embeddings

+
+
+ +
+
+ +
+
@@ -62,4 +50,61 @@ head %}
+
+

Protected phrases

+
+
+ +
+
+ +
+
+ +
+
+ +
+
+ + + + + + + + + + + + + + + + + + + + + + + +
CandidatesJudgedUnjudgedProtectedBooks indexedBooks generatedBooks fully judged
{{ phrase_stats.candidate_phrases }}{{ phrase_stats.judged_candidates }}{{ phrase_stats.unjudged_candidates }}{{ phrase_stats.protected_phrases }}{{ phrase_stats.total_books }}{{ phrase_stats.books_with_candidates }}{{ phrase_stats.books_fully_judged }}
+
{% endblock %} diff --git a/python/ebook_search/api/templates/book_detail.html b/python/ebook_search/api/templates/book_detail.html index 1102bec..54fbdcb 100644 --- a/python/ebook_search/api/templates/book_detail.html +++ b/python/ebook_search/api/templates/book_detail.html @@ -30,6 +30,15 @@ > +
+ +

Candidate n-grams

diff --git a/python/ebook_search/bm25_corpus.py b/python/ebook_search/bm25_corpus.py index 0450d64..b0612be 100644 --- a/python/ebook_search/bm25_corpus.py +++ b/python/ebook_search/bm25_corpus.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import json import logging import shutil @@ -17,7 +18,7 @@ from sqlalchemy import func, select, union_all from python.orm.richie import EbookChapter, EbookChunk, EbookSource if TYPE_CHECKING: - from sqlalchemy.orm import Session + from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.config import EbookSearchConfig @@ -73,14 +74,14 @@ def get_current_bm25_index(index_path: Path) -> Path: return index_path -def ensure_bm25_corpus(session: Session, config: EbookSearchConfig) -> None: +async def ensure_bm25_corpus(session: AsyncSession, config: EbookSearchConfig) -> None: """Create or refresh the persisted BM25 corpus when it is missing or stale.""" index_path = bm25_index_path(config) manifest = read_bm25_manifest(index_path) - db_updated_at = corpus_last_updated_at(session) + db_updated_at = await corpus_last_updated_at(session) if not bm25_index_exists(index_path, manifest): logger.info("ebook_bm25_index_missing path=%s", index_path) - refresh_bm25_corpus(session, config, db_updated_at=db_updated_at) + await refresh_bm25_corpus(session, config, db_updated_at=db_updated_at) return if db_updated_at is not None and manifest is not None and manifest.created_at < db_updated_at: logger.info( @@ -89,7 +90,7 @@ def ensure_bm25_corpus(session: Session, config: EbookSearchConfig) -> None: manifest.created_at.isoformat(), db_updated_at.isoformat(), ) - refresh_bm25_corpus(session, config, db_updated_at=db_updated_at) + await refresh_bm25_corpus(session, config, db_updated_at=db_updated_at) return logger.info( "ebook_bm25_index_current path=%s chunks=%s created_at=%s", @@ -99,21 +100,24 @@ def ensure_bm25_corpus(session: Session, config: EbookSearchConfig) -> None: ) -def refresh_bm25_corpus( - session: Session, +async def refresh_bm25_corpus( + session: AsyncSession, config: EbookSearchConfig, *, db_updated_at: datetime | None = None, ) -> BM25Manifest: - """Rebuild and persist the BM25 corpus from the current database chunks.""" + """Rebuild and persist the BM25 corpus from the current database chunks. + + The index build is CPU and disk work, so it runs in a worker thread. + """ index_path = bm25_index_path(config) - records, texts = fetch_bm25_corpus_records(session) + records, texts = await fetch_bm25_corpus_records(session) manifest = BM25Manifest( created_at=datetime.now(tz=UTC), - db_updated_at=db_updated_at if db_updated_at is not None else corpus_last_updated_at(session), + db_updated_at=db_updated_at if db_updated_at is not None else await corpus_last_updated_at(session), chunk_count=len(records), ) - write_bm25_corpus(index_path, records, texts, manifest) + await asyncio.to_thread(write_bm25_corpus, index_path, records, texts, manifest) logger.info( "ebook_bm25_index_refreshed path=%s chunks=%s created_at=%s", index_path, @@ -164,7 +168,7 @@ def score_bm25_corpus(query: str, corpus: BM25Corpus, *, limit: int) -> list[tup return results -def fetch_bm25_corpus_records(session: Session) -> tuple[list[dict[str, object]], list[str]]: +async def fetch_bm25_corpus_records(session: AsyncSession) -> tuple[list[dict[str, object]], list[str]]: """Fetch persistable BM25 corpus records and their matching index texts from the database. search_text is only needed to build the index, so it is returned separately instead of @@ -188,21 +192,21 @@ def fetch_bm25_corpus_records(session: Session) -> tuple[list[dict[str, object]] ) records: list[dict[str, object]] = [] texts: list[str] = [] - for row in session.execute(statement).mappings(): + for row in (await session.execute(statement)).mappings(): record = dict(row) texts.append(str(record.pop("bm25_text"))) records.append(record) return records, texts -def corpus_last_updated_at(session: Session) -> datetime | None: +async def corpus_last_updated_at(session: AsyncSession) -> datetime | None: """Return the latest source/chapter/chunk update timestamp relevant to BM25 text.""" update_times = union_all( select(func.max(EbookSource.updated).label("updated")), select(func.max(EbookChapter.updated).label("updated")), select(func.max(EbookChunk.updated).label("updated")), ).subquery() - return session.scalar(select(func.max(update_times.c.updated))) + return await session.scalar(select(func.max(update_times.c.updated))) def write_bm25_corpus( diff --git a/python/ebook_search/config.py b/python/ebook_search/config.py index 173f07b..687d705 100644 --- a/python/ebook_search/config.py +++ b/python/ebook_search/config.py @@ -88,6 +88,9 @@ class EbookSearchConfig(BaseSettings): bm25_refresh_delay_seconds: int = 60 protected_phrase_max_candidates_per_book: int = 5000 protected_phrase_llm_candidates_per_book: int = 500 + protected_phrase_extraction_workers: int = 16 + phrase_judge_book_workers: int = 20 + phrase_judge_phrase_workers: int = 100 protected_phrase_confidence_threshold: float = 0.80 phrase_matching_enabled: bool = True phrase_hit_boost: float = 0.25 diff --git a/python/ebook_search/embeddings.py b/python/ebook_search/embeddings.py index f542e2b..003136e 100644 --- a/python/ebook_search/embeddings.py +++ b/python/ebook_search/embeddings.py @@ -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) diff --git a/python/ebook_search/ingest.py b/python/ebook_search/ingest.py index b5f9254..51fc801 100644 --- a/python/ebook_search/ingest.py +++ b/python/ebook_search/ingest.py @@ -2,6 +2,7 @@ from __future__ import annotations +import asyncio import hashlib import logging from dataclasses import dataclass @@ -21,7 +22,7 @@ DEFAULT_CHUNK_TOKENS = 700 DEFAULT_CHUNK_OVERLAP = 100 if TYPE_CHECKING: - from sqlalchemy.orm import Session + from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.config import EbookSearchConfig from python.ebook_search.epub_parse import ParsedChapter @@ -73,45 +74,63 @@ def chunk_text( return [chunk for chunk in chunks if chunk.text] -def ingest_configured_paths(session: Session, config: EbookSearchConfig) -> int: +def find_library_epubs(library_path: str) -> tuple[Path, list[Path] | None]: + """Resolve one configured library path and collect its EPUB files (blocking filesystem walk). + + Returns: + tuple[Path, list[Path] | None]: The expanded path and its EPUB files, or ``None`` when + the path is neither an EPUB file nor a directory. + """ + path = Path(library_path).expanduser() + if path.is_file() and path.suffix.lower() == ".epub": + return path, [path] + if path.is_dir(): + return path, sorted(path.rglob("*.epub")) + return path, None + + +async def ingest_configured_paths(session: AsyncSession, config: EbookSearchConfig) -> int: """Ingest every EPUB found under configured library paths.""" count = 0 for library_path in config.library_paths: - path = Path(library_path).expanduser() + path, epub_paths = await asyncio.to_thread(find_library_epubs, library_path) logger.info("ebook_ingest_path_start path=%s", path) - if path.is_file() and path.suffix.lower() == ".epub": - count += int(ingest_file(session, path, config)) - elif path.is_dir(): - for epub_path in sorted(path.rglob("*.epub")): - count += int(ingest_file(session, epub_path, config)) - else: + if epub_paths is None: logger.warning("ebook_ingest_path_missing path=%s", path) + continue + for epub_path in epub_paths: + count += int(await ingest_file(session, epub_path, config)) logger.info("ebook_ingest_paths_complete changed_files=%s configured_paths=%s", count, len(config.library_paths)) return count -def ingest_file(session: Session, path: Path, config: EbookSearchConfig) -> bool: +def resolve_ingest_path(path: Path) -> Path: + """Expand and resolve an ingest path (blocking filesystem call).""" + return path.expanduser().resolve() + + +async def ingest_file(session: AsyncSession, path: Path, config: EbookSearchConfig) -> bool: """Ingest one EPUB file. Return True when the database changed.""" try: - resolved_path = path.expanduser().resolve() + resolved_path = await asyncio.to_thread(resolve_ingest_path, path) logger.info("ebook_ingest_file_start path=%s", resolved_path) - file_hash = sha256_file(resolved_path) - existing = find_existing_source(session, resolved_path, file_hash) + file_hash = await asyncio.to_thread(sha256_file, resolved_path) + existing = await find_existing_source(session, resolved_path, file_hash) if existing is not None and existing.file_sha256 == file_hash: stat = resolved_path.stat() existing.file_path = str(resolved_path) existing.file_mtime = datetime.fromtimestamp(stat.st_mtime, tz=UTC) existing.file_size = stat.st_size - session.flush() + await session.flush() logger.info("ebook_ingest_file_unchanged source_id=%s path=%s", existing.id, resolved_path) return False if existing is not None: logger.info("ebook_ingest_file_replacing source_id=%s path=%s", existing.id, resolved_path) - session.delete(existing) - session.flush() + await session.delete(existing) + await session.flush() stat = resolved_path.stat() - parsed = parse_epub(resolved_path) + parsed = await asyncio.to_thread(parse_epub, resolved_path) source = EbookSource( title=parsed.title, author=parsed.author, @@ -124,7 +143,7 @@ def ingest_file(session: Session, path: Path, config: EbookSearchConfig) -> bool file_size=stat.st_size, ) session.add(source) - session.flush() + await session.flush() chunk_index = 0 for spine_index, parsed_chapter in enumerate(parsed.chapters): @@ -135,11 +154,11 @@ def ingest_file(session: Session, path: Path, config: EbookSearchConfig) -> bool href=parsed_chapter.href, ) session.add(chapter) - session.flush() + await session.flush() chunk_index = add_chapter_chunks(session, source, chapter, parsed_chapter, chunk_index, config) - session.commit() - mention_count = index_chunk_phrase_mentions_for_book(session, source.id, config) + await session.commit() + mention_count = await index_chunk_phrase_mentions_for_book(session, source.id, config) logger.info( "ebook_ingest_file_complete source_id=%s path=%s chapters=%s chunks=%s phrase_mentions=%s", source.id, @@ -155,15 +174,15 @@ def ingest_file(session: Session, path: Path, config: EbookSearchConfig) -> bool return True -def find_existing_source(session: Session, path: Path, file_hash: str) -> EbookSource | None: +async def find_existing_source(session: AsyncSession, path: Path, file_hash: str) -> EbookSource | None: """Find an existing source by canonical path or file hash.""" - return session.scalar( + return await session.scalar( select(EbookSource).where(or_(EbookSource.file_path == str(path), EbookSource.file_sha256 == file_hash)) ) def add_chapter_chunks( - session: Session, + session: AsyncSession, source: EbookSource, chapter: EbookChapter, parsed_chapter: ParsedChapter, diff --git a/python/ebook_search/llm_interface.py b/python/ebook_search/llm_interface.py index e573630..0412ee6 100644 --- a/python/ebook_search/llm_interface.py +++ b/python/ebook_search/llm_interface.py @@ -22,10 +22,26 @@ def auth_headers(api_key: str) -> dict[str, str]: return {"Authorization": f"Bearer {api_key}"} -def request_embeddings(texts: Sequence[str], config: EbookSearchConfig) -> list[list[float]]: - """Request embeddings from the configured OpenAI-compatible endpoint.""" +async def request_embeddings( + client: httpx.AsyncClient, + texts: Sequence[str], + config: EbookSearchConfig, +) -> list[list[float]]: + """Request embeddings from the configured OpenAI-compatible endpoint. + + Args: + client (httpx.AsyncClient): Shared async client for LLM calls. + texts (Sequence[str]): Texts to embed. + config (EbookSearchConfig): Runtime settings supplying the endpoint, model, and auth. + + Returns: + list[list[float]]: One embedding vector per input text. + + Raises: + RuntimeError: If the request fails or the response cannot be parsed. + """ try: - response = httpx.post( + response = await client.post( f"{config.embedding_base_url.rstrip('/')}/embeddings", headers=auth_headers(config.embedding_api_key), json={"model": config.embedding_model, "input": list(texts)}, @@ -44,10 +60,15 @@ def request_embeddings(texts: Sequence[str], config: EbookSearchConfig) -> list[ raise RuntimeError(msg) from error -def check_embedding_endpoint(config: EbookSearchConfig, *, timeout_seconds: float = 5.0) -> bool: +async def check_embedding_endpoint( + client: httpx.AsyncClient, + config: EbookSearchConfig, + *, + timeout_seconds: float = 5.0, +) -> bool: """Return whether the configured embedding endpoint answers a model listing.""" try: - response = httpx.get( + response = await client.get( f"{config.embedding_base_url.rstrip('/')}/models", headers=auth_headers(config.embedding_api_key), timeout=timeout_seconds, @@ -59,10 +80,15 @@ def check_embedding_endpoint(config: EbookSearchConfig, *, timeout_seconds: floa return True -def check_chat_endpoint(config: EbookSearchConfig, *, timeout_seconds: float = 5.0) -> bool: +async def check_chat_endpoint( + client: httpx.AsyncClient, + config: EbookSearchConfig, + *, + timeout_seconds: float = 5.0, +) -> bool: """Return whether the configured chat (answering) endpoint answers a model listing.""" try: - response = httpx.get( + response = await client.get( f"{config.vllm_base_url.rstrip('/')}/models", headers=auth_headers(config.vllm_api_key), timeout=timeout_seconds, @@ -98,18 +124,29 @@ def embedding_vectors_from_response(body: object) -> list[list[float]]: return vectors -def request_rerank( +async def request_rerank( + client: httpx.AsyncClient, query: str, documents: Sequence[str], config: RerankConfig, ) -> object | None: - """Request rerank scores from the configured vLLM endpoint.""" + """Request rerank scores from the configured vLLM endpoint. + + Args: + client (httpx.AsyncClient): Shared async client for LLM calls. + query (str): Query the documents are scored against. + documents (Sequence[str]): Candidate documents to score. + config (RerankConfig): Rerank endpoint settings. + + Returns: + object | None: The decoded response body, or ``None`` when it is not valid JSON. + """ payload = { "model": config.model, "query": query, "documents": list(documents), } - response = httpx.post( + response = await client.post( f"{config.base_url.rstrip('/')}/rerank", json=payload, timeout=config.timeout_seconds, @@ -122,13 +159,26 @@ def request_rerank( return None -def request_chat_completion( +async def request_chat_completion( + client: httpx.AsyncClient, config: EbookSearchConfig, messages: Sequence[dict[str, str]], ) -> str: - """Request a chat completion from the configured OpenAI-compatible endpoint.""" + """Request a chat completion over a shared async client. + + Args: + client (httpx.AsyncClient): Shared async client whose connection pool bounds concurrency. + config (EbookSearchConfig): Runtime settings supplying the endpoint, model, and auth. + messages (Sequence[dict[str, str]]): OpenAI-style chat messages. + + Returns: + str: The assistant message text. + + Raises: + RuntimeError: If the request fails or the response cannot be parsed. + """ try: - response = httpx.post( + response = await client.post( f"{config.vllm_base_url.rstrip('/')}/chat/completions", headers=auth_headers(config.vllm_api_key), json={ diff --git a/python/ebook_search/protected_phrases/extraction.py b/python/ebook_search/protected_phrases/extraction.py index 3af80a5..1370e5b 100644 --- a/python/ebook_search/protected_phrases/extraction.py +++ b/python/ebook_search/protected_phrases/extraction.py @@ -4,7 +4,7 @@ from __future__ import annotations import logging import re -from collections import defaultdict +from collections import Counter, defaultdict from functools import lru_cache from time import perf_counter from typing import TYPE_CHECKING, Protocol @@ -15,6 +15,7 @@ from python.ebook_search.protected_phrases.config import ( get_bad_ends, get_bad_starts, get_ignored_phrases, + get_junk_tokens, get_most_common_words, ) from python.ebook_search.protected_phrases.models import PhraseCandidate @@ -28,16 +29,10 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -BAD_START_SCORE_PENALTY = 3.0 -BAD_END_SCORE_PENALTY = 3.0 -SOURCE_FIELDS = ( - "source_raw_ngram", - "source_yake", - "source_spacy_ner", - "source_spacy_noun_chunk", - "source_capitalized", - "source_metadata", -) +BAD_START_SCORE_PENALTY = 10.0 +BAD_END_SCORE_PENALTY = 10.0 +MULTI_SOURCE_SCORE_BONUS = 2.0 +MULTI_SOURCE_MIN_SOURCES = 2 CAPITALIZED_PHRASE_RE = re.compile(r"\b(?:[A-Z][a-zA-Z']+)(?:\s+(?:of|the|and|in|on|for|[A-Z][a-zA-Z']+)){0,6}") @@ -91,21 +86,6 @@ class YakeExtractorFactory(Protocol): """ -def strip_leading_articles(phrase_norm: str) -> str: - """Remove one leading English article from a normalized phrase. - - Args: - phrase_norm (str): Normalized phrase text to strip. - - Returns: - str: The phrase with a single leading ``the``, ``a``, or ``an`` removed. - """ - tokens_ = phrase_norm.split() - if tokens_ and tokens_[0] in {"the", "a", "an"}: - tokens_ = tokens_[1:] - return " ".join(tokens_) - - def normalize_candidate_phrase( phrase_text: str, config: EbookSearchConfig, @@ -146,35 +126,71 @@ def normalize_candidate_phrase( return display_text or phrase_norm, phrase_norm, len(selected_tokens) -def extract_raw_ngrams(text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]: - """Extract raw normalized n-grams as high-recall candidates. +def count_raw_ngrams(tokens: Sequence[str], config: EbookSearchConfig) -> Counter[str]: + """Count every n-gram window in one normalized token block. + + ``tokens`` are already normalized (see :func:`tokenize`), so each window's normalized form + is the joined tokens directly. Counting into a plain :class:`Counter` rather than + :class:`PhraseCandidate` objects keeps this hot loop cheap; callers filter ignored phrases + and materialize candidates per unique phrase afterwards, which is far fewer operations than + doing either per window. Args: - text (str): Book text to slide n-gram windows over. + tokens (Sequence[str]): Normalized tokens for one text block. config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: - dict[str, PhraseCandidate]: Candidates keyed by normalized phrase, with raw counts. + Counter[str]: Raw occurrence counts keyed by normalized phrase. """ - tokens_ = tokenize(text) - out: dict[str, PhraseCandidate] = {} - for ngram_size in range(config.phrase_min_tokens, config.phrase_max_tokens + 1): - for start in range(len(tokens_) - ngram_size + 1): - normalized = normalize_candidate_phrase(" ".join(tokens_[start : start + ngram_size]), config) - if normalized is None: - continue - phrase_text, phrase_norm, token_count = normalized - item = out.setdefault( - phrase_norm, - PhraseCandidate( - phrase_text=phrase_text, - phrase_norm=phrase_norm, - token_count=token_count, - source_raw_ngram=True, - ), - ) - item.raw_count += 1 - return out + return Counter( + " ".join(tokens[start : start + ngram_size]) + for ngram_size in range(config.phrase_min_tokens, config.phrase_max_tokens + 1) + for start in range(len(tokens) - ngram_size + 1) + ) + + +def extract_raw_ngrams_by_chapter( + chapters: Sequence[str], + config: EbookSearchConfig, +) -> dict[str, PhraseCandidate]: + """Extract raw n-grams across chapters, tracking both raw counts and chapter spread. + + Counting each chapter separately makes chapter spread fall out of dict membership: a phrase's + ``chapter_count`` is simply how many per-chapter count maps contain it, so no per-window seen + tracking is needed. This also lets the enrichment step skip re-sliding the same n-gram sizes. + + Phrases below the minimum raw count are dropped here rather than materialized: most unique + n-grams occur once, and :func:`filter_storable_candidates` would discard them as too rare + anyway, so building ``PhraseCandidate`` objects for them is wasted work. + + Args: + chapters (Sequence[str]): Chapter-like text blocks to slide n-gram windows over. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + dict[str, PhraseCandidate]: Candidates meeting the minimum raw count, keyed by normalized + phrase, with raw and chapter counts. + """ + chapter_count_maps = [count_raw_ngrams(tokenize(chapter), config) for chapter in chapters] + total_counts: Counter[str] = Counter() + chapter_spread: Counter[str] = Counter() + for chapter_counts in chapter_count_maps: + total_counts.update(chapter_counts) + chapter_spread.update(chapter_counts.keys()) + min_raw_count = minimum_candidate_raw_count(config) + ignored = get_ignored_phrases() + return { + phrase_norm: PhraseCandidate( + phrase_text=phrase_norm, + phrase_norm=phrase_norm, + token_count=phrase_norm.count(" ") + 1, + source_raw_ngram=True, + raw_count=raw_count, + chapter_count=chapter_spread[phrase_norm], + ) + for phrase_norm, raw_count in total_counts.items() + if raw_count >= min_raw_count and phrase_norm not in ignored + } @lru_cache(maxsize=2) @@ -381,6 +397,7 @@ def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None: existing.source_capitalized = existing.source_capitalized or item.source_capitalized existing.source_metadata = existing.source_metadata or item.source_metadata existing.raw_count += item.raw_count + existing.chapter_count = max(existing.chapter_count, item.chapter_count) if item.yake_score is not None: existing.yake_score = item.yake_score if item.spacy_label: @@ -390,12 +407,19 @@ def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None: def enrich_with_frequency_and_chapter_counts( candidates: Mapping[str, PhraseCandidate], chapters: Sequence[str], + *, + counted_sizes: Iterable[int] = (), ) -> dict[str, PhraseCandidate]: """Add raw occurrence and chapter-spread counts to candidates. + Candidates whose ``token_count`` is in ``counted_sizes`` are left untouched: those counts + were already computed while sliding the chapters in :func:`extract_raw_ngrams_by_chapter`, + so re-sliding those n-gram sizes here would just duplicate that work. + Args: candidates (Mapping[str, PhraseCandidate]): Candidates to enrich, keyed by normalized phrase. chapters (Sequence[str]): Chapter-like text blocks used to count occurrences and spread. + counted_sizes (Iterable[int]): Token counts whose counts are already populated and should be skipped. Returns: dict[str, PhraseCandidate]: Candidates with updated ``raw_count`` and ``chapter_count`` values. @@ -403,10 +427,40 @@ def enrich_with_frequency_and_chapter_counts( if not candidates: return {} + already_counted = set(counted_sizes) candidate_sets_by_size: dict[int, set[str]] = defaultdict(set) for phrase_norm, candidate in candidates.items(): + if candidate.token_count in already_counted: + continue candidate_sets_by_size[candidate.token_count].add(phrase_norm) + enriched = dict(candidates) + if not candidate_sets_by_size: + return enriched + + total_counts, chapter_counts = count_candidate_occurrences(candidate_sets_by_size, chapters) + for phrase_norm, candidate in enriched.items(): + if candidate.token_count in already_counted: + continue + candidate.raw_count = max(candidate.raw_count, total_counts[phrase_norm]) + candidate.chapter_count = chapter_counts[phrase_norm] + return enriched + + +def count_candidate_occurrences( + candidate_sets_by_size: Mapping[int, set[str]], + chapters: Sequence[str], +) -> tuple[dict[str, int], dict[str, int]]: + """Count total occurrences and chapter spread for candidate phrases across chapters. + + Args: + candidate_sets_by_size (Mapping[int, set[str]]): Candidate normalized phrases grouped by token count. + chapters (Sequence[str]): Chapter-like text blocks to slide n-gram windows over. + + Returns: + tuple[dict[str, int], dict[str, int]]: Total occurrence counts and chapter-spread counts, + each keyed by normalized phrase. + """ total_counts: defaultdict[str, int] = defaultdict(int) chapter_counts: defaultdict[str, int] = defaultdict(int) for chapter in chapters: @@ -421,18 +475,13 @@ def enrich_with_frequency_and_chapter_counts( seen_in_chapter.add(phrase_norm) for phrase_norm in seen_in_chapter: chapter_counts[phrase_norm] += 1 - - enriched = dict(candidates) - for phrase_norm, candidate in enriched.items(): - candidate.raw_count = max(candidate.raw_count, total_counts[phrase_norm]) - candidate.chapter_count = chapter_counts[phrase_norm] - return enriched + return total_counts, chapter_counts def filter_storable_candidates( candidates: Mapping[str, PhraseCandidate], config: EbookSearchConfig, -) -> tuple[dict[str, PhraseCandidate], int, int, int]: +) -> tuple[dict[str, PhraseCandidate], int, int, int, int]: """Remove candidates that should not be persisted. Args: @@ -440,14 +489,15 @@ def filter_storable_candidates( config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: - tuple[dict[str, PhraseCandidate], int, int, int]: The storable candidates followed by the counts - dropped for being too short, too rare, and too common. + tuple[dict[str, PhraseCandidate], int, int, int, int]: The storable candidates followed by the + counts dropped for being too short, too rare, too common, and junk. """ min_raw_count = minimum_candidate_raw_count(config) filtered: dict[str, PhraseCandidate] = {} too_short = 0 too_rare = 0 too_common = 0 + junk = 0 for phrase_norm, candidate in candidates.items(): if candidate.token_count < config.phrase_min_tokens: too_short += 1 @@ -455,11 +505,15 @@ def filter_storable_candidates( if candidate.raw_count < min_raw_count: too_rare += 1 continue - if is_most_common_word_phrase(phrase_norm): + phrase_tokens = phrase_norm.split() + if is_most_common_word_phrase(phrase_tokens): too_common += 1 continue + if is_junk_phrase(phrase_tokens): + junk += 1 + continue filtered[phrase_norm] = candidate - return filtered, too_short, too_rare, too_common + return filtered, too_short, too_rare, too_common, junk def minimum_candidate_raw_count(config: EbookSearchConfig) -> int: @@ -474,18 +528,41 @@ def minimum_candidate_raw_count(config: EbookSearchConfig) -> int: return max(config.phrase_raw_ngram_min_count, 1) -def is_most_common_word_phrase(phrase_norm: str) -> bool: +def is_most_common_word_phrase(phrase_tokens: list[str]) -> bool: """Return whether every token in a normalized phrase is a common word. Args: - phrase_norm (str): Normalized phrase text to inspect. + phrase_tokens (list[str]): Normalized phrase tokens to inspect. Returns: bool: True when the phrase is non-empty and every token is a common word. """ - tokens_ = phrase_norm.split() common_words = get_most_common_words() - return bool(tokens_) and all(token in common_words for token in tokens_) + return bool(phrase_tokens) and all(token in common_words for token in phrase_tokens) + + +def is_junk_phrase(phrase_tokens: list[str]) -> bool: + """Return whether a normalized phrase is lexical junk not worth LLM judging. + + Judged data shows phrases containing a dialogue/action verb or a pronoun contraction are + never kept, and phrases whose tokens are mostly common words almost never are. Possessives + of proper nouns (``chapman's death``) pass because matching is by exact token, and + exactly-half-common bigrams (``data feed``) pass because the common-word rule is strict. + + Args: + phrase_tokens (list[str]): Normalized phrase tokens to inspect. + + Returns: + bool: True when the phrase contains a junk token or is majority common words. + """ + if not phrase_tokens: + return False + junk_tokens = get_junk_tokens() + if any(token in junk_tokens for token in phrase_tokens): + return True + common_words = get_most_common_words() + half_phrase_len = len(phrase_tokens) // 2 + return sum(token in common_words for token in phrase_tokens) > half_phrase_len def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> float: @@ -499,6 +576,8 @@ def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> fl float: Combined score from sources, frequency, and length, less any penalties. """ score = source_score(candidate) + frequency_score(candidate, config) + token_count_score(candidate, config) + if non_raw_source_count(candidate) >= MULTI_SOURCE_MIN_SOURCES: + score += MULTI_SOURCE_SCORE_BONUS if candidate.phrase_norm in get_ignored_phrases(): score -= 100.0 if has_bad_start(candidate.phrase_norm): @@ -508,6 +587,26 @@ def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> fl return score +def non_raw_source_count(candidate: PhraseCandidate) -> int: + """Count the non-raw-ngram extraction sources that produced a candidate. + + Args: + candidate (PhraseCandidate): Candidate whose enabled sources are counted. + + Returns: + int: Number of enabled sources other than the raw n-gram slide. + """ + return sum( + ( + candidate.source_yake, + candidate.source_spacy_ner, + candidate.source_spacy_noun_chunk, + candidate.source_capitalized, + candidate.source_metadata, + ) + ) + + def has_bad_start(phrase_norm: str) -> bool: """Return whether a normalized phrase starts with a bad starting token. @@ -517,8 +616,8 @@ def has_bad_start(phrase_norm: str) -> bool: Returns: bool: True when the first token is a known bad starting token. """ - tokens_ = phrase_norm.split() - return bool(tokens_ and tokens_[0] in get_bad_starts()) + phrase_tokens = phrase_norm.split() + return bool(phrase_tokens and phrase_tokens[0] in get_bad_starts()) def has_bad_end(phrase_norm: str) -> bool: @@ -530,8 +629,8 @@ def has_bad_end(phrase_norm: str) -> bool: Returns: bool: True when the last token is a known bad ending token. """ - tokens_ = phrase_norm.split() - return bool(tokens_ and tokens_[-1] in get_bad_ends()) + phrase_tokens = phrase_norm.split() + return bool(phrase_tokens and phrase_tokens[-1] in get_bad_ends()) def source_score(candidate: PhraseCandidate) -> float: @@ -550,6 +649,7 @@ def source_score(candidate: PhraseCandidate) -> float: (candidate.source_spacy_ner, 2.5), (candidate.source_spacy_noun_chunk, 1.5), (candidate.source_capitalized, 2.0), + (candidate.source_metadata, 2.0), (candidate.source_raw_ngram, 0.5), ) if enabled @@ -569,10 +669,10 @@ def frequency_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> fl return sum( weight for count, threshold, weight in ( - (candidate.raw_count, config.phrase_raw_count_score_threshold, 1.0), - (candidate.raw_count, config.phrase_raw_count_high_score_threshold, 1.0), - (candidate.chapter_count, config.phrase_chapter_count_score_threshold, 1.0), - (candidate.chapter_count, config.phrase_chapter_count_high_score_threshold, 1.0), + (candidate.raw_count, config.phrase_raw_count_score_threshold, 0.5), + (candidate.raw_count, config.phrase_raw_count_high_score_threshold, 0.5), + (candidate.chapter_count, config.phrase_chapter_count_score_threshold, 0.5), + (candidate.chapter_count, config.phrase_chapter_count_high_score_threshold, 0.5), ) if count >= threshold ) @@ -679,7 +779,7 @@ def extract_phrase_candidates_for_book( config.protected_phrase_max_candidates_per_book, ) raw_started_at = perf_counter() - raw = extract_raw_ngrams(book_text, config) + raw = extract_raw_ngrams_by_chapter(chapters, config) logger.info( "ebook_phrase_candidate_extract_raw_complete candidates=%s duration_ms=%.1f", len(raw), @@ -713,11 +813,16 @@ def extract_phrase_candidates_for_book( candidates = merge_candidate_sources(raw, yake_candidates, spacy_candidates, capitalized, metadata_candidates) enriched_started_at = perf_counter() - candidates = enrich_with_frequency_and_chapter_counts(candidates, chapters) - pre_filter_count = len(candidates) - candidates, filtered_too_short, filtered_too_rare, filtered_too_common = filter_storable_candidates( + # Raw n-gram sizes were already counted per chapter above, so only enrich the remaining + # (entity-length) sizes here instead of re-sliding every size over the whole book. + candidates = enrich_with_frequency_and_chapter_counts( candidates, - config, + chapters, + counted_sizes=range(config.phrase_min_tokens, config.phrase_max_tokens + 1), + ) + pre_filter_count = len(candidates) + candidates, filtered_too_short, filtered_too_rare, filtered_too_common, filtered_junk = filter_storable_candidates( + candidates, config ) for candidate in candidates.values(): candidate.candidate_score = score_candidate(candidate, config) @@ -727,8 +832,8 @@ def extract_phrase_candidates_for_book( ] logger.info( "ebook_phrase_candidate_extract_complete raw=%s yake=%s spacy=%s capitalized=%s metadata=%s " - "merged=%s filtered_too_short=%s filtered_too_rare=%s filtered_too_common=%s min_uses=%s " - "storable=%s limited=%s enrich_score_ms=%.1f duration_ms=%.1f", + "merged=%s filtered_too_short=%s filtered_too_rare=%s filtered_too_common=%s filtered_junk=%s " + "min_uses=%s storable=%s limited=%s enrich_score_ms=%.1f duration_ms=%.1f", len(raw), len(yake_candidates), len(spacy_candidates), @@ -738,6 +843,7 @@ def extract_phrase_candidates_for_book( filtered_too_short, filtered_too_rare, filtered_too_common, + filtered_junk, minimum_candidate_raw_count(config), len(candidates), len(limited), diff --git a/python/ebook_search/protected_phrases/generate_ngrams.py b/python/ebook_search/protected_phrases/generate_ngrams.py index c3a6499..165d9b5 100644 --- a/python/ebook_search/protected_phrases/generate_ngrams.py +++ b/python/ebook_search/protected_phrases/generate_ngrams.py @@ -2,7 +2,9 @@ from __future__ import annotations +import asyncio import logging +from collections import deque from time import perf_counter from typing import TYPE_CHECKING @@ -14,44 +16,55 @@ from python.ebook_search.protected_phrases.models import ( PhraseCandidateGenerationResult, PhraseRecalculationResult, ) +from python.ebook_search.protected_phrases.pool import extract_phrase_candidates_in_pool, get_extraction_pool from python.ebook_search.protected_phrases.store import ( + bulk_upsert_unjudged_candidates, delete_phrase_data_for_book, load_book_chapter_texts, metadata_for_source, + new_candidate_row, prune_unstorable_unjudged_candidate_phrases, - save_candidate_to_db, ) -from python.orm.richie import EbookSource +from python.orm.richie import EbookCandidatePhrase, EbookSource if TYPE_CHECKING: from collections.abc import Mapping, Sequence + from concurrent.futures import Future - from sqlalchemy.orm import Session + from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.config import EbookSearchConfig from python.ebook_search.protected_phrases.extraction import SpacyLanguage - from python.orm.richie import EbookCandidatePhrase + from python.ebook_search.protected_phrases.models import PhraseCandidate logger = logging.getLogger(__name__) -def generate_candidate_phrases_for_books( - session: Session, +async def generate_candidate_phrases_for_books( + session: AsyncSession, config: EbookSearchConfig, *, - nlp: SpacyLanguage | None = None, + only_missing: bool = False, ) -> PhraseCandidateGenerationResult: """Create or refresh candidate phrases for indexed books without calling the LLM judge. + Extraction always runs concurrently in the shared process pool so a full backfill uses + multiple cores. + Args: session (Session): Active database session. config (EbookSearchConfig): Runtime phrase-tuning settings. - nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. + only_missing (bool): When True, only generate for books that have no candidate phrases + yet instead of refreshing every book. Returns: PhraseCandidateGenerationResult: Per-corpus counts of books seen, built, and candidates stored. """ - sources = session.scalars(select(EbookSource).order_by(EbookSource.id)).all() + source_query = select(EbookSource).order_by(EbookSource.id) + if only_missing: + has_candidates = select(EbookCandidatePhrase.id).where(EbookCandidatePhrase.book_id == EbookSource.id) + source_query = source_query.where(~has_candidates.exists()) + sources = (await session.scalars(source_query)).all() books_seen = len(sources) logger.info( "ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s max_candidates_per_book=%s", @@ -61,7 +74,7 @@ def generate_candidate_phrases_for_books( config.protected_phrase_max_candidates_per_book, ) - outcomes = [generate_candidates_for_source(session, source, config, nlp=nlp) for source in sources] + outcomes = await generate_candidates_for_sources_pooled(session, sources, config) result = PhraseCandidateGenerationResult( books_seen=books_seen, @@ -77,76 +90,99 @@ def generate_candidate_phrases_for_books( return result -def generate_candidates_for_source( - session: Session, - source: EbookSource, +async def generate_candidates_for_sources_pooled( + session: AsyncSession, + sources: Sequence[EbookSource], config: EbookSearchConfig, - *, - nlp: SpacyLanguage | None = None, -) -> BookCandidateResult: - """Generate and store candidate phrases for one book, managing its own transaction. +) -> list[BookCandidateResult]: + """Generate candidate phrases for many books, extracting them concurrently in worker processes. - Commits on success; rolls back and re-raises on error so callers stop the backfill. + Chapter loading and row persistence stay on the caller's session (serial), while the CPU-bound + extraction runs in the shared process pool. A bounded window of in-flight books overlaps + extraction across cores without loading every book's candidates into memory at once. Args: session (Session): Active database session. - source (EbookSource): Indexed book to generate candidates for. + sources (Sequence[EbookSource]): Indexed books to generate candidates for. config (EbookSearchConfig): Runtime phrase-tuning settings. - nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. Returns: - BookCandidateResult: Candidate count and whether the book was committed. + list[BookCandidateResult]: One result per book. + """ + pool = get_extraction_pool(config.protected_phrase_extraction_workers) + max_in_flight = max(1, config.protected_phrase_extraction_workers) * 2 + pending: deque[tuple[EbookSource, Future[list[PhraseCandidate]]]] = deque() + outcomes: list[BookCandidateResult] = [] + + async def drain_one() -> None: + source, future = pending.popleft() + extracted = await asyncio.wrap_future(future) + outcomes.append(await store_source_candidates(session, source, extracted, config)) + + try: + for source in sources: + chapters = await load_book_chapter_texts(session, source.id) + if not chapters: + logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id) + outcomes.append(BookCandidateResult()) + continue + future = pool.submit( + extract_phrase_candidates_for_book, + "\n\n".join(chapters), + chapters, + config, + metadata=metadata_for_source(source), + ) + pending.append((source, future)) + if len(pending) >= max_in_flight: + await drain_one() + while pending: + await drain_one() + except Exception: + for _, future in pending: + future.cancel() + await session.rollback() + logger.exception("ebook_candidate_phrase_generation_pooled_failed") + raise + return outcomes + + +async def store_source_candidates( + session: AsyncSession, + source: EbookSource, + limited_candidates: list[PhraseCandidate], + config: EbookSearchConfig, +) -> BookCandidateResult: + """Persist and commit one book's already-extracted candidates. + + Args: + session (AsyncSession): Active database session. + source (EbookSource): Book the candidates belong to. + limited_candidates (list[PhraseCandidate]): Scored candidates to persist. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + BookCandidateResult: Candidate count and that the book was committed. """ book_started_at = perf_counter() - logger.info( - "ebook_candidate_phrase_generation_book_start source_id=%s title=%r", - source.id, - source.title, - ) - try: - chapters = load_book_chapter_texts(session, source.id) - if not chapters: - logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id) - return BookCandidateResult() - book_text = "\n\n".join(chapters) - logger.info( - "ebook_candidate_phrase_generation_book_loaded source_id=%s chapters=%s chars=%s", - source.id, - len(chapters), - len(book_text), - ) - - candidates = generate_candidate_phrases_for_book( - session, - source.id, - series_id=None, - book_text=book_text, - chapters=chapters, - config=config, - nlp=nlp, - metadata=metadata_for_source(source), - ) - session.commit() - except Exception: - session.rollback() - logger.exception("ebook_candidate_phrase_generation_book_failed source_id=%s", source.id) - raise - + saved_count = await store_candidate_phrases_for_book(session, source.id, None, limited_candidates, config) + await session.commit() logger.info( "ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s duration_ms=%.1f", source.id, - len(candidates), + saved_count, (perf_counter() - book_started_at) * 1000, ) - return BookCandidateResult(candidates=len(candidates), built=True) + return BookCandidateResult(candidates=saved_count, built=True) -def recalculate_candidate_phrases_for_book( - session: Session, +async def recalculate_candidate_phrases_for_book( + session: AsyncSession, source: EbookSource, config: EbookSearchConfig, *, nlp: SpacyLanguage | None = None, + use_process_pool: bool = False, ) -> PhraseRecalculationResult: """Remove all book phrase data, regenerate candidates, and commit the completed book. @@ -155,6 +191,9 @@ def recalculate_candidate_phrases_for_book( source (EbookSource): Indexed book to recalculate. config (EbookSearchConfig): Runtime phrase-tuning settings. nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. + use_process_pool (bool): Run the CPU-bound extraction in a worker process so concurrent + recalculations do not serialize behind the GIL. Defaults to in-process for callers + (tests, backfills) that do not need it. Returns: PhraseRecalculationResult: Deleted-row counts and the number of candidates regenerated. @@ -166,11 +205,11 @@ def recalculate_candidate_phrases_for_book( source.title, ) try: - deleted = delete_phrase_data_for_book(session, source.id) - chapters = load_book_chapter_texts(session, source.id) + deleted = await delete_phrase_data_for_book(session, source.id) + chapters = await load_book_chapter_texts(session, source.id) if not chapters: logger.warning("ebook_candidate_phrase_recalculation_book_empty source_id=%s", source.id) - session.commit() + await session.commit() return PhraseRecalculationResult( book_id=source.id, deleted_candidates=deleted.deleted_candidates, @@ -180,19 +219,20 @@ def recalculate_candidate_phrases_for_book( candidate_phrases=0, ) - candidates = generate_candidate_phrases_for_book( + candidate_count = await generate_candidate_phrases_for_book( session, source.id, series_id=None, - book_text="\n\n".join(chapters), chapters=chapters, config=config, nlp=nlp, metadata=metadata_for_source(source), + replace_all=True, + use_process_pool=use_process_pool, ) - session.commit() + await session.commit() except Exception: - session.rollback() + await session.rollback() logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id) raise @@ -202,7 +242,7 @@ def recalculate_candidate_phrases_for_book( deleted_protected_phrases=deleted.deleted_protected_phrases, deleted_aliases=deleted.deleted_aliases, deleted_mentions=deleted.deleted_mentions, - candidate_phrases=len(candidates), + candidate_phrases=candidate_count, ) logger.info( "ebook_candidate_phrase_recalculation_complete source_id=%s deleted_candidates=%s " @@ -218,57 +258,113 @@ def recalculate_candidate_phrases_for_book( return result -def generate_candidate_phrases_for_book( - session: Session, +async def generate_candidate_phrases_for_book( + session: AsyncSession, book_id: int, series_id: int | None, - book_text: str, chapters: Sequence[str], config: EbookSearchConfig, *, nlp: SpacyLanguage | None = None, metadata: Mapping[str, object] | None = None, -) -> list[EbookCandidatePhrase]: + replace_all: bool = False, + use_process_pool: bool = False, +) -> int: """Extract and store candidate phrases for one book without LLM judging. Args: session (Session): Active database session. book_id (int): Book the candidates belong to. series_id (int | None): Series scope for the stored candidates. - book_text (str): Full book text used for extraction. - chapters (Sequence[str]): Chapter-like text blocks used for frequency counts. + chapters (Sequence[str]): Chapter-like text blocks used for extraction and frequency counts. config (EbookSearchConfig): Runtime phrase-tuning settings. nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source. + replace_all (bool): When the caller has already cleared this book's candidates (e.g. a + recalculation), skip the per-candidate existence lookup and bulk-insert new rows. + use_process_pool (bool): Run the CPU-bound extraction in a worker process to avoid + serializing concurrent requests behind the GIL. Ignored when ``nlp`` is set, since + the spaCy pipeline cannot be sent to a worker process. Returns: - list[EbookCandidatePhrase]: The stored candidate phrase rows. + int: Number of candidate phrase rows stored. """ started_at = perf_counter() - limited_candidates = extract_phrase_candidates_for_book( - book_text, - chapters, + book_text = "\n\n".join(chapters) + if use_process_pool and nlp is None: + limited_candidates = await extract_phrase_candidates_in_pool(book_text, chapters, config, metadata=metadata) + else: + limited_candidates = extract_phrase_candidates_for_book( + book_text, + chapters, + config, + nlp=nlp, + metadata=metadata, + ) + saved_count = await store_candidate_phrases_for_book( + session, + book_id, + series_id, + limited_candidates, config, - nlp=nlp, - metadata=metadata, + replace_all=replace_all, ) - save_started_at = perf_counter() - pruned_count = prune_unstorable_unjudged_candidate_phrases(session, book_id, config) logger.info( - "ebook_candidate_phrase_save_start book_id=%s candidates=%s pruned_unstorable=%s", + "ebook_candidate_phrase_generation_book_duration book_id=%s candidates=%s duration_ms=%.1f", book_id, - len(limited_candidates), - pruned_count, - ) - rows = [ - save_candidate_to_db(session, book_id, series_id, candidate, judgment=None) for candidate in limited_candidates - ] - session.flush() - logger.info( - "ebook_candidate_phrase_generation_complete book_id=%s candidates=%s save_ms=%.1f duration_ms=%.1f", - book_id, - len(rows), - (perf_counter() - save_started_at) * 1000, + saved_count, (perf_counter() - started_at) * 1000, ) - return rows + return saved_count + + +async def store_candidate_phrases_for_book( + session: AsyncSession, + book_id: int, + series_id: int | None, + limited_candidates: list[PhraseCandidate], + config: EbookSearchConfig, + *, + replace_all: bool = False, +) -> int: + """Persist already-extracted candidate phrase rows for one book without committing. + + Args: + session (Session): Active database session. + book_id (int): Book the candidates belong to. + series_id (int | None): Series scope for the stored candidates. + limited_candidates (list[PhraseCandidate]): Scored candidates to persist. + config (EbookSearchConfig): Runtime phrase-tuning settings. + replace_all (bool): When the caller has already cleared this book's candidates, skip the + per-candidate existence lookup and bulk-insert new rows. + + Returns: + int: Number of candidate phrase rows stored. + """ + save_started_at = perf_counter() + if replace_all: + rows = [new_candidate_row(book_id, series_id, candidate) for candidate in limited_candidates] + session.add_all(rows) + await session.flush() + saved_count = len(rows) + logger.info( + "ebook_candidate_phrase_save_start book_id=%s candidates=%s mode=bulk_insert", + book_id, + len(limited_candidates), + ) + else: + pruned_count = await prune_unstorable_unjudged_candidate_phrases(session, book_id, config) + logger.info( + "ebook_candidate_phrase_save_start book_id=%s candidates=%s pruned_unstorable=%s", + book_id, + len(limited_candidates), + pruned_count, + ) + saved_count = await bulk_upsert_unjudged_candidates(session, book_id, series_id, limited_candidates) + logger.info( + "ebook_candidate_phrase_save_complete book_id=%s candidates=%s save_ms=%.1f", + book_id, + saved_count, + (perf_counter() - save_started_at) * 1000, + ) + return saved_count diff --git a/python/ebook_search/protected_phrases/judge_ngrams.py b/python/ebook_search/protected_phrases/judge_ngrams.py index 1be6e17..56b6b70 100644 --- a/python/ebook_search/protected_phrases/judge_ngrams.py +++ b/python/ebook_search/protected_phrases/judge_ngrams.py @@ -2,20 +2,24 @@ from __future__ import annotations +import asyncio import json import logging import re from time import perf_counter from typing import TYPE_CHECKING +import httpx from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.llm_interface import request_chat_completion from python.ebook_search.protected_phrases.extraction import ( candidate_source_names, get_sample_contexts, + is_junk_phrase, is_most_common_word_phrase, - minimum_candidate_raw_count, + score_candidate, ) from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book from python.ebook_search.protected_phrases.models import BookJudgmentResult, LLMJudgment, PhraseJudgmentBackfillResult @@ -32,47 +36,65 @@ from python.ebook_search.protected_phrases.text_normalization import normalize_t from python.orm.richie import EbookSource if TYPE_CHECKING: - from sqlalchemy.orm import Session + from collections.abc import Sequence + + from sqlalchemy.ext.asyncio import AsyncEngine from python.ebook_search.config import EbookSearchConfig from python.ebook_search.protected_phrases.models import PhraseCandidate - from python.orm.richie import EbookCandidatePhrase, EbookProtectedPhrase + from python.orm.richie import EbookProtectedPhrase JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL) - logger = logging.getLogger(__name__) -def judge_candidate_phrases_for_books( - session: Session, +async def judge_candidate_phrases_for_books( + engine: AsyncEngine, config: EbookSearchConfig, + *, + source_ids: Sequence[int] | None = None, ) -> PhraseJudgmentBackfillResult: - """Judge stored candidate phrases and promote accepted phrases for indexed books. + """Judge candidate phrases for books, fanning LLM calls out across books and phrases. + + Up to ``phrase_judge_book_workers`` books are judged at once, and within each book candidates + are judged in concurrent chunks of ``phrase_judge_phrase_workers``. Each book uses its own + short-lived sessions for reads and writes; no database connection is held while LLM calls are + in flight. For a pseudo-single-threaded run (solo testing, debugging), set both worker + settings to 1. Args: - session (Session): Active database session. - config (EbookSearchConfig): Runtime phrase-tuning settings. + engine (AsyncEngine): Engine used to open one session per book. + config (EbookSearchConfig): Runtime phrase-tuning settings and chat configuration. + source_ids (Sequence[int] | None): Books to judge; ``None`` judges every indexed book. Returns: PhraseJudgmentBackfillResult: Per-corpus counts of books judged, failures, candidates, protected phrases, and mentions. """ - source_ids = session.scalars(select(EbookSource.id).order_by(EbookSource.id)).all() + if source_ids is None: + async with AsyncSession(engine) as session: + source_ids = list((await session.scalars(select(EbookSource.id).order_by(EbookSource.id))).all()) books_seen = len(source_ids) + book_workers = max(1, config.phrase_judge_book_workers) + phrase_workers = max(1, config.phrase_judge_phrase_workers) logger.info( - "ebook_candidate_phrase_judgment_start books_seen=%s llm_candidates_per_book=%s " - "target_protected_per_book=%s confidence_threshold=%.2f min_tokens=%s min_uses=%s", + "ebook_candidate_phrase_judgment_start books_seen=%s book_workers=%s phrase_workers=%s " + "confidence_threshold=%.2f", books_seen, - config.protected_phrase_llm_candidates_per_book, - config.phrase_target_protected_per_book, + book_workers, + phrase_workers, config.protected_phrase_confidence_threshold, - config.phrase_min_tokens, - minimum_candidate_raw_count(config), ) - outcomes = [judge_book_for_backfill(session, source_id, config) for source_id in source_ids] + book_semaphore = asyncio.Semaphore(book_workers) + max_connections = book_workers * phrase_workers + limits = httpx.Limits(max_connections=max_connections, max_keepalive_connections=max_connections) + async with httpx.AsyncClient(limits=limits) as client: + outcomes = await asyncio.gather( + *(judge_one_book_async(engine, source_id, config, client, book_semaphore) for source_id in source_ids) + ) result = PhraseJudgmentBackfillResult( books_seen=books_seen, @@ -95,235 +117,247 @@ def judge_candidate_phrases_for_books( return result -def judge_book_for_backfill( - session: Session, +async def judge_one_book_async( + engine: AsyncEngine, source_id: int, config: EbookSearchConfig, + client: httpx.AsyncClient, + book_semaphore: asyncio.Semaphore, ) -> BookJudgmentResult: - """Judge one book's candidates and return its outcome. + """Judge one book concurrently and persist the outcome, honoring the book-level limit. Args: - session (Session): Active database session. + engine (AsyncEngine): Engine used to open the book's read and write sessions. source_id (int): Book to judge candidates for. config (EbookSearchConfig): Runtime phrase-tuning settings. + client (httpx.AsyncClient): Shared async client for LLM calls. + book_semaphore (asyncio.Semaphore): Caps how many books judge at once. Returns: - BookJudgmentResult: The book's judgment outcome, or an empty result when nothing was unjudged. + BookJudgmentResult: The book's judgment outcome. """ - unjudged_count = count_unjudged_candidates(session, source_id, config) - if not unjudged_count: - logger.info( - "ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s", - source_id, - ) - return BookJudgmentResult() - logger.info( - "ebook_candidate_phrase_judgment_book_start source_id=%s unjudged=%s", - source_id, - unjudged_count, - ) - return run_book_judgment(session, source_id, unjudged_count, config) + async with book_semaphore: + try: + prepared = await prepare_book_judgment(engine, source_id, config) + if prepared is None: + return BookJudgmentResult() + work_items, target_remaining = prepared + judged = await judge_book_candidates_async(client, config, source_id, work_items, target_remaining) + if not judged: + return BookJudgmentResult() + return await persist_book_judgments(engine, source_id, config, judged) + except Exception: + logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id) + return BookJudgmentResult(failed=True) -def run_book_judgment( - session: Session, +async def prepare_book_judgment( + engine: AsyncEngine, source_id: int, - unjudged_count: int, config: EbookSearchConfig, -) -> BookJudgmentResult: - """Judge and index one book's candidates, managing its own transaction. - - Commits on success and rolls back on error, returning a :class:`BookJudgmentResult` - that describes the outcome rather than raising it to the caller. +) -> tuple[list[tuple[int, PhraseCandidate]], int | None] | None: + """Load one book's candidates to judge, with sample contexts, on a short-lived read session. Args: - session (Session): Active database session. - source_id (int): Book to judge candidates for. - unjudged_count (int): Storable unjudged candidates counted before judging, for logging. + engine (AsyncEngine): Engine used to open the read session. + source_id (int): Book to load candidates for. config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: - BookJudgmentResult: Judged, protected, and mention counts with commit and failure flags. - """ - book_started_at = perf_counter() - try: - book_text = load_book_text(session, source_id) - if not book_text: - logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id) - return BookJudgmentResult() - - judged, protected = judge_candidate_phrases_for_book( - session, - source_id, - series_id=None, - normalized_book_text=normalize_text(book_text), - config=config, - ) - if judged == 0: - logger.info( - "ebook_candidate_phrase_judgment_book_skip_no_judgments source_id=%s unjudged=%s", - source_id, - unjudged_count, - ) - return BookJudgmentResult() - - mentions = index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0 - session.commit() - except Exception: - session.rollback() - logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id) - return BookJudgmentResult(failed=True) - - logger.info( - "ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s " - "duration_ms=%.1f", - source_id, - judged, - len(protected), - mentions, - (perf_counter() - book_started_at) * 1000, - ) - return BookJudgmentResult(judged=judged, protected=len(protected), mentions=mentions, committed=True) - - -def judge_candidate_phrases_for_book( - session: Session, - book_id: int, - series_id: int | None, - normalized_book_text: str, - config: EbookSearchConfig, -) -> tuple[int, list[EbookProtectedPhrase]]: - """Judge unjudged candidate phrase rows for one book. - - Args: - session (Session): Active database session. - book_id (int): Book whose candidates are judged. - series_id (int | None): Series scope for promoted protected phrases. - normalized_book_text (str): Whole book text, already normalized, for context lookups. - config (EbookSearchConfig): Runtime phrase-tuning settings. - - Returns: - tuple[int, list[EbookProtectedPhrase]]: Number of candidates judged and the promoted phrases. + tuple[list[tuple[int, PhraseCandidate]], int | None] | None: Candidate rows paired with + in-memory candidates and the remaining protected-phrase target, or ``None`` when the book + has nothing to judge. """ judgment_limit = config.protected_phrase_llm_candidates_per_book if judgment_limit <= 0: - logger.info("ebook_candidate_phrase_judgment_skipped_llm_limit_zero book_id=%s", book_id) - return 0, [] - - existing_protected = count_protected_phrases(session, book_id) - target_remaining: int | None = None - if config.phrase_target_protected_per_book > 0: - target_remaining = max(config.phrase_target_protected_per_book - existing_protected, 0) - if target_remaining == 0: - logger.info( - "ebook_candidate_phrase_judgment_skipped_target_met book_id=%s existing_protected=%s target=%s", - book_id, - existing_protected, - config.phrase_target_protected_per_book, + return None + async with AsyncSession(engine) as session: + if not await count_unjudged_candidates(session, source_id, config): + logger.info("ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s", source_id) + return None + existing_protected = await count_protected_phrases(session, source_id) + target_remaining: int | None = None + if config.phrase_target_protected_per_book > 0: + target_remaining = max(config.phrase_target_protected_per_book - existing_protected, 0) + if target_remaining == 0: + logger.info( + "ebook_candidate_phrase_judgment_skipped_target_met source_id=%s existing_protected=%s target=%s", + source_id, + existing_protected, + config.phrase_target_protected_per_book, + ) + return None + book_text = await load_book_text(session, source_id) + if not book_text: + logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id) + return None + normalized_book_text = normalize_text(book_text) + # Stored rows may predate the current junk filters and score weights, so re-filter and + # rescore every unjudged row here instead of trusting the persisted candidate_score. + rows = await load_candidates_for_judgment(session, source_id, config) + scored_items: list[tuple[int, PhraseCandidate]] = [] + skipped_junk = 0 + for row in rows: + candidate = phrase_candidate_from_row(row) + if is_junk_phrase(candidate.phrase_norm.split()): + skipped_junk += 1 + continue + candidate.candidate_score = score_candidate(candidate, config) + scored_items.append((row.id, candidate)) + scored_items.sort(key=lambda item: item[1].candidate_score, reverse=True) + work_items = scored_items[:judgment_limit] + for _, candidate in work_items: + candidate.sample_contexts = candidate.sample_contexts or get_sample_contexts( + normalized_book_text, candidate.phrase_norm ) - return 0, [] - - rows = load_candidates_for_judgment(session, book_id, judgment_limit, config) logger.info( - "ebook_candidate_phrase_judgment_candidates_loaded book_id=%s candidates=%s existing_protected=%s " - "target_remaining=%s judgment_limit=%s", - book_id, + "ebook_candidate_phrase_judgment_candidates_loaded source_id=%s candidates=%s skipped_junk=%s " + "unjudged_rows=%s existing_protected=%s target_remaining=%s judgment_limit=%s", + source_id, + len(work_items), + skipped_junk, len(rows), existing_protected, target_remaining, judgment_limit, ) - - judged_count = 0 - protected: list[EbookProtectedPhrase] = [] - for row_number, row in enumerate(rows, start=1): - candidate, judgment, candidate_row = judge_candidate_row( - session, book_id, series_id, normalized_book_text, row, row_number, len(rows), config - ) - judged_count += 1 - if not should_protect_judged_candidate(row, candidate, judgment, book_id, config): - continue - protected.append(upsert_protected_phrase(session, book_id, series_id, candidate, judgment, candidate_row)) - if target_remaining is not None and len(protected) >= target_remaining: - break - - session.flush() - return judged_count, protected + return work_items, target_remaining -def judge_candidate_row( - session: Session, - book_id: int, - series_id: int | None, - normalized_book_text: str, - row: EbookCandidatePhrase, - row_number: int, - total_rows: int, +async def judge_book_candidates_async( + client: httpx.AsyncClient, config: EbookSearchConfig, -) -> tuple[PhraseCandidate, LLMJudgment, EbookCandidatePhrase]: - """Run and persist the LLM judgment for a single candidate row. + source_id: int, + work_items: list[tuple[int, PhraseCandidate]], + target_remaining: int | None, +) -> list[tuple[int, PhraseCandidate, LLMJudgment, bool]]: + """Judge a book's candidates in concurrent chunks, stopping once the target is reached. + + Promotion decisions are made in memory so judging can stop early without any database writes. Args: - session (Session): Active database session. - book_id (int): Book the candidate belongs to. - series_id (int | None): Series scope for the saved candidate. - normalized_book_text (str): Whole book text, already normalized, for context lookups. - row (EbookCandidatePhrase): Stored candidate row to judge. - row_number (int): 1-based position of the row in the batch, for logging. - total_rows (int): Total rows in the batch, for logging. + client (httpx.AsyncClient): Shared async client for LLM calls. config (EbookSearchConfig): Runtime phrase-tuning settings. + source_id (int): Book being judged, for logging. + work_items (list[tuple[int, PhraseCandidate]]): Candidate row ids paired with candidates, + in best-first score order. + target_remaining (int | None): Remaining protected-phrase target, or ``None`` for no cap. Returns: - tuple[PhraseCandidate, LLMJudgment, EbookCandidatePhrase]: The candidate, its judgment, - and the persisted candidate row. + list[tuple[int, PhraseCandidate, LLMJudgment, bool]]: Judged rows with their judgment and + whether each should be promoted. """ - row_started_at = perf_counter() - candidate = phrase_candidate_from_row(row) - candidate.sample_contexts = row.sample_contexts or get_sample_contexts(normalized_book_text, candidate.phrase_norm) + chunk_size = max(1, config.phrase_judge_phrase_workers) + judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]] = [] + promoted = 0 + for start in range(0, len(work_items), chunk_size): + chunk = work_items[start : start + chunk_size] + judgments = await asyncio.gather(*(judge_candidate_async(client, config, candidate) for _, candidate in chunk)) + for (candidate_id, candidate), judgment in zip(chunk, judgments, strict=True): + promote = (target_remaining is None or promoted < target_remaining) and should_protect_judged_candidate( + candidate, judgment, source_id, config, candidate_id=candidate_id + ) + if promote: + promoted += 1 + judged.append((candidate_id, candidate, judgment, promote)) + if target_remaining is not None and promoted >= target_remaining: + break + return judged + + +async def judge_candidate_async( + client: httpx.AsyncClient, + config: EbookSearchConfig, + candidate: PhraseCandidate, +) -> LLMJudgment: + """Judge one candidate with the LLM over the shared async client. + + Args: + client (httpx.AsyncClient): Shared async client for LLM calls. + config (EbookSearchConfig): Runtime phrase-tuning settings. + candidate (PhraseCandidate): Candidate to judge. + + Returns: + LLMJudgment: The parsed judgment. + """ + content = await request_chat_completion(client, config, build_judge_messages(candidate)) + return parse_llm_judgment(content, config) + + +async def persist_book_judgments( + engine: AsyncEngine, + source_id: int, + config: EbookSearchConfig, + judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]], +) -> BookJudgmentResult: + """Persist one book's judgments and promotions in a single committed transaction. + + Args: + engine (AsyncEngine): Engine used to open the write session. + source_id (int): Book being persisted. + config (EbookSearchConfig): Runtime phrase-tuning settings. + judged (list[tuple[int, PhraseCandidate, LLMJudgment, bool]]): Judged candidates with their + judgment and promotion flag. + + Returns: + BookJudgmentResult: The book's committed counts, or a failed result on error. + """ + book_started_at = perf_counter() + async with AsyncSession(engine, expire_on_commit=False) as session: + try: + protected: list[EbookProtectedPhrase] = [] + for candidate_id, candidate, judgment, promote in judged: + candidate_row = await save_candidate_to_db(session, source_id, None, candidate, judgment=judgment) + if promote: + protected.append( + await upsert_protected_phrase(session, source_id, None, candidate, judgment, candidate_row) + ) + logger.info( + "ebook_candidate_phrase_judgment_candidate_complete source_id=%s candidate_id=%s phrase=%r " + "keep=%s confidence=%.3f category=%r promoted=%s", + source_id, + candidate_id, + candidate.phrase_norm, + judgment.keep, + judgment.confidence, + judgment.category, + promote, + ) + await session.flush() + mentions = await index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0 + await session.commit() + except Exception: + await session.rollback() + logger.exception("ebook_candidate_phrase_judgment_book_persist_failed source_id=%s", source_id) + return BookJudgmentResult(failed=True) logger.info( - "ebook_candidate_phrase_judgment_candidate_start book_id=%s candidate_id=%s row_number=%s rows=%s " - "phrase=%r score=%.3f raw_count=%s chapter_count=%s", - book_id, - row.id, - row_number, - total_rows, - candidate.phrase_norm, - candidate.candidate_score, - candidate.raw_count, - candidate.chapter_count, + "ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s " + "duration_ms=%.1f", + source_id, + len(judged), + len(protected), + mentions, + (perf_counter() - book_started_at) * 1000, ) - judgment = judge_candidate_with_llm(candidate, config) - candidate_row = save_candidate_to_db(session, book_id, series_id, candidate, judgment=judgment) - logger.info( - "ebook_candidate_phrase_judgment_candidate_complete book_id=%s candidate_id=%s phrase=%r keep=%s " - "confidence=%.3f importance=%.3f category=%r duration_ms=%.1f", - book_id, - row.id, - candidate.phrase_norm, - judgment.keep, - judgment.confidence, - judgment.importance, - judgment.category, - (perf_counter() - row_started_at) * 1000, - ) - return candidate, judgment, candidate_row + return BookJudgmentResult(judged=len(judged), protected=len(protected), mentions=mentions, committed=True) def should_protect_judged_candidate( - row: EbookCandidatePhrase, candidate: PhraseCandidate, judgment: LLMJudgment, book_id: int, config: EbookSearchConfig, + *, + candidate_id: int, ) -> bool: """Report whether a judged candidate qualifies to become a protected phrase. Args: - row (EbookCandidatePhrase): Stored candidate row the judgment came from. candidate (PhraseCandidate): In-memory candidate that was judged. judgment (LLMJudgment): Judge decision for the candidate. book_id (int): Book the candidate belongs to, for logging. config (EbookSearchConfig): Runtime phrase-tuning settings. + candidate_id (int): Stored candidate row id the judgment came from, for logging. Returns: bool: True when the judged candidate should be promoted to a protected phrase. @@ -331,25 +365,26 @@ def should_protect_judged_candidate( if not judgment.keep or judgment.confidence < config.protected_phrase_confidence_threshold: return False accepted_norm = normalize_text(judgment.canonical or candidate.phrase_text) - accepted_token_count = len(accepted_norm.split()) + accepted_tokens = accepted_norm.split() + accepted_token_count = len(accepted_tokens) if accepted_token_count < config.phrase_min_tokens: logger.info( "ebook_candidate_phrase_judgment_candidate_skip_short_canonical book_id=%s candidate_id=%s " "phrase=%r canonical=%r token_count=%s min_tokens=%s", book_id, - row.id, + candidate_id, candidate.phrase_norm, accepted_norm, accepted_token_count, config.phrase_min_tokens, ) return False - if is_most_common_word_phrase(accepted_norm): + if is_most_common_word_phrase(accepted_tokens): logger.info( "ebook_candidate_phrase_judgment_candidate_skip_common_canonical book_id=%s candidate_id=%s " "phrase=%r canonical=%r", book_id, - row.id, + candidate_id, candidate.phrase_norm, accepted_norm, ) @@ -357,18 +392,14 @@ def should_protect_judged_candidate( return True -"""LLM judging of extracted candidate phrases.""" - - -def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConfig) -> LLMJudgment: - """Ask the configured chat model to judge one pre-extracted candidate. +def build_judge_messages(candidate: PhraseCandidate) -> list[dict[str, str]]: + """Build the chat messages used to judge one candidate phrase. Args: - candidate (PhraseCandidate): Candidate to send to the LLM judge. - config (EbookSearchConfig): Runtime phrase-tuning settings and chat configuration. + candidate (PhraseCandidate): Candidate to describe for the judge. Returns: - LLMJudgment: The parsed structured judgment for the candidate. + list[dict[str, str]]: OpenAI-style system and user messages. """ payload = { "phrase": candidate.phrase_norm, @@ -378,7 +409,7 @@ def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConf "chapter_count": candidate.chapter_count, "contexts": candidate.sample_contexts, } - messages = [ + return [ { "role": "system", "content": ( @@ -392,7 +423,6 @@ def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConf }, {"role": "user", "content": json.dumps(payload, ensure_ascii=True)}, ] - return parse_llm_judgment(request_chat_completion(config, messages), config) def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment: diff --git a/python/ebook_search/protected_phrases/matching.py b/python/ebook_search/protected_phrases/matching.py index 7fecbe3..172a852 100644 --- a/python/ebook_search/protected_phrases/matching.py +++ b/python/ebook_search/protected_phrases/matching.py @@ -26,7 +26,7 @@ from python.orm.richie import ( if TYPE_CHECKING: from collections.abc import Iterator, Sequence - from sqlalchemy.orm import Session + from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.config import EbookSearchConfig from python.ebook_search.protected_phrases.text_normalization import NormalizedToken @@ -34,8 +34,8 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def load_phrase_lookup( - session: Session, +async def load_phrase_lookup( + session: AsyncSession, config: EbookSearchConfig, *, book_id: int | None = None, @@ -44,7 +44,7 @@ def load_phrase_lookup( """Load protected phrases and aliases into RAM lookup maps. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. config (EbookSearchConfig): Runtime phrase-tuning settings. book_id (int | None): Optional book scope to restrict loaded phrases. series_id (int | None): Optional series scope to restrict loaded phrases. @@ -65,7 +65,7 @@ def load_phrase_lookup( if scope_filter is not None: phrase_statement = phrase_statement.where(scope_filter) - for row in session.execute(phrase_statement): + for row in await session.execute(phrase_statement): phrase_id = int(row.id) phrase_norm = str(row.phrase_norm) norm_to_ids[phrase_norm].append(phrase_id) @@ -78,7 +78,7 @@ def load_phrase_lookup( if scope_filter is not None: alias_statement = alias_statement.where(scope_filter) - for row in session.execute(alias_statement): + for row in await session.execute(alias_statement): alias_norm = str(row.alias_norm) alias_to_ids[alias_norm].append(int(row.phrase_id)) max_tokens = max(max_tokens, len(alias_norm.split())) @@ -197,11 +197,11 @@ def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], loo return matches -def hydrate_matches(session: Session, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]: +async def hydrate_matches(session: AsyncSession, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]: """Fetch protected phrase metadata for raw phrase matches. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. matches (Sequence[PhraseMatch]): Unhydrated matches to enrich. Returns: @@ -216,7 +216,7 @@ def hydrate_matches(session: Session, matches: Sequence[PhraseMatch]) -> list[Hy rows = { row.id: row - for row in session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids))) + for row in await session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids))) } hydrated: list[HydratedPhraseMatch] = [] for match in matches: @@ -332,8 +332,8 @@ def resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhr return kept -def detect_protected_phrases_for_query( - session: Session, +async def detect_protected_phrases_for_query( + session: AsyncSession, query_text: str, config: EbookSearchConfig, *, @@ -344,7 +344,7 @@ def detect_protected_phrases_for_query( """Run the full online protected-phrase query-detection pipeline. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. query_text (str): User query text to detect phrases in. config (EbookSearchConfig): Runtime phrase-tuning settings. lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``. @@ -355,13 +355,15 @@ def detect_protected_phrases_for_query( list[HydratedPhraseMatch]: Hydrated, overlap-resolved phrase matches for the query. """ active_lookup = ( - lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id) + lookup + if lookup is not None + else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id) ) - return resolve_overlaps(hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup))) + return resolve_overlaps(await hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup))) -def index_chunk_phrase_mentions_for_book( - session: Session, +async def index_chunk_phrase_mentions_for_book( + session: AsyncSession, book_id: int, config: EbookSearchConfig, *, @@ -371,7 +373,7 @@ def index_chunk_phrase_mentions_for_book( """Rebuild chunk phrase mentions for all chunks in one book. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. book_id (int): Book whose chunk mentions are rebuilt. config (EbookSearchConfig): Runtime phrase-tuning settings. series_id (int | None): Optional series scope for lookup loading. @@ -381,23 +383,25 @@ def index_chunk_phrase_mentions_for_book( int: Total number of chunk phrase mentions indexed for the book. """ active_lookup = ( - lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id) + lookup + if lookup is not None + else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id) ) - session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id)) - chunks = session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id)) + await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id)) + chunks = await session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id)) count = 0 for chunk in chunks: - count += index_chunk_phrase_mentions(session, chunk, lookup=active_lookup) - session.flush() + count += await index_chunk_phrase_mentions(session, chunk, lookup=active_lookup) + await session.flush() logger.info("ebook_chunk_phrase_mentions_indexed book_id=%s mentions=%s", book_id, count) return count -def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup: PhraseLookup) -> int: +async def index_chunk_phrase_mentions(session: AsyncSession, chunk: EbookChunk, *, lookup: PhraseLookup) -> int: """Store protected phrase mentions for one chunk. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. chunk (EbookChunk): Chunk whose text is scanned for phrase mentions. lookup (PhraseLookup): In-memory phrase and alias lookup maps. @@ -405,7 +409,7 @@ def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup: int: Number of phrase mentions stored for the chunk. """ raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup) - hydrated = resolve_overlaps(hydrate_matches(session, raw_matches)) + hydrated = resolve_overlaps(await hydrate_matches(session, raw_matches)) for match in hydrated: session.add( EbookChunkPhraseMention( @@ -420,8 +424,8 @@ def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup: return len(hydrated) -def phrase_hits_for_chunks( - session: Session, +async def phrase_hits_for_chunks( + session: AsyncSession, *, chunk_ids: Sequence[int], phrase_ids: Sequence[int], @@ -429,7 +433,7 @@ def phrase_hits_for_chunks( """Return matched protected phrases with mention counts by chunk id using indexed chunk mentions. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. chunk_ids (Sequence[int]): Chunk ids to look up mentions for. phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to. @@ -456,7 +460,7 @@ def phrase_hits_for_chunks( .order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text) ) hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list) - for row in session.execute(statement): + for row in await session.execute(statement): hits[int(row.chunk_id)].append( ChunkPhraseHit( phrase_id=int(row.phrase_id), @@ -467,8 +471,8 @@ def phrase_hits_for_chunks( return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()} -def phrase_hit_counts_for_chunks( - session: Session, +async def phrase_hit_counts_for_chunks( + session: AsyncSession, *, chunk_ids: Sequence[int], phrase_ids: Sequence[int], @@ -476,12 +480,12 @@ def phrase_hit_counts_for_chunks( """Return phrase-hit counts by chunk id using indexed chunk mentions. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. chunk_ids (Sequence[int]): Chunk ids to count mentions for. phrase_ids (Sequence[int]): Protected phrase ids to restrict the counts to. Returns: dict[int, int]: Total mention count per chunk id. """ - hits = phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids) + hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids) return {chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in hits.items()} diff --git a/python/ebook_search/protected_phrases/models.py b/python/ebook_search/protected_phrases/models.py index 2336fe6..36f810f 100644 --- a/python/ebook_search/protected_phrases/models.py +++ b/python/ebook_search/protected_phrases/models.py @@ -188,6 +188,29 @@ class PhraseCandidateGenerationResult: candidate_phrases: int +@dataclass(frozen=True, slots=True) +class CorpusPhraseStats: + """Corpus-wide candidate and protected phrase counts for the admin page. + + Attributes: + total_books (int): Indexed books in the corpus. + books_with_candidates (int): Books that have candidate phrases generated. + books_fully_judged (int): Books with candidates where every candidate has been judged. + candidate_phrases (int): Candidate phrases stored across all books. + judged_candidates (int): Candidate phrases that have been LLM judged. + unjudged_candidates (int): Candidate phrases still waiting for judgment. + protected_phrases (int): Protected phrases promoted across all books. + """ + + total_books: int + books_with_candidates: int + books_fully_judged: int + candidate_phrases: int + judged_candidates: int + unjudged_candidates: int + protected_phrases: int + + @dataclass(frozen=True, slots=True) class PhraseJudgmentBackfillResult: """Summary of LLM judging for stored candidate phrases. diff --git a/python/ebook_search/protected_phrases/pool.py b/python/ebook_search/protected_phrases/pool.py new file mode 100644 index 0000000..68fdec4 --- /dev/null +++ b/python/ebook_search/protected_phrases/pool.py @@ -0,0 +1,101 @@ +"""Process pool for offloading CPU-bound phrase extraction off the request thread. + +Phrase extraction is pure-Python CPU work (n-gram sliding, YAKE), so running it inline in a +sync request handler serializes concurrent recalculations behind the GIL. Submitting it to a +``ProcessPoolExecutor`` lets concurrent extractions run in parallel across cores instead. A +``spawn`` context is used so workers do not inherit the parent's database engine, connections, +or server threads. +""" + +from __future__ import annotations + +import asyncio +import logging +import multiprocessing +import os +from concurrent.futures import ProcessPoolExecutor +from threading import Lock +from typing import TYPE_CHECKING + +from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from python.ebook_search.config import EbookSearchConfig + from python.ebook_search.protected_phrases.models import PhraseCandidate + +logger = logging.getLogger(__name__) + + +class _ExtractionPool: + """Lazily created process-wide extraction pool and the lock guarding it.""" + + def __init__(self) -> None: + self.lock = Lock() + self.pool: ProcessPoolExecutor | None = None + + +_extraction_pool = _ExtractionPool() + + +def get_extraction_pool(max_workers: int) -> ProcessPoolExecutor: + """Return the shared extraction process pool, creating it on first use. + + Args: + max_workers (int): Desired worker count; values below 1 fall back to the CPU count. + + Returns: + ProcessPoolExecutor: The shared pool for phrase extraction. + """ + with _extraction_pool.lock: + if _extraction_pool.pool is None: + workers = max_workers if max_workers > 0 else (os.cpu_count() or 1) + _extraction_pool.pool = ProcessPoolExecutor( + max_workers=workers, + mp_context=multiprocessing.get_context("spawn"), + ) + logger.info("ebook_phrase_extraction_pool_started workers=%s", workers) + return _extraction_pool.pool + + +def shutdown_extraction_pool() -> None: + """Shut down the shared extraction pool if it was started.""" + with _extraction_pool.lock: + if _extraction_pool.pool is not None: + _extraction_pool.pool.shutdown(wait=False, cancel_futures=True) + _extraction_pool.pool = None + logger.info("ebook_phrase_extraction_pool_shutdown") + + +async def extract_phrase_candidates_in_pool( + book_text: str, + chapters: Sequence[str], + config: EbookSearchConfig, + *, + metadata: Mapping[str, object] | None, +) -> list[PhraseCandidate]: + """Run book phrase extraction in a worker process and await the result. + + Only the CPU-bound extraction runs in the worker; the caller keeps all database work in the + request process. The spaCy pipeline is not supported here because it is not picklable, so + this always runs the non-spaCy extraction path. + + Args: + book_text (str): Full book text used for extraction. + chapters (Sequence[str]): Chapter-like text blocks used for frequency counts. + config (EbookSearchConfig): Runtime phrase-tuning settings. + metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source. + + Returns: + list[PhraseCandidate]: Scored candidates sorted best-first and capped per book. + """ + pool = get_extraction_pool(config.protected_phrase_extraction_workers) + future = pool.submit( + extract_phrase_candidates_for_book, + book_text, + list(chapters), + config, + metadata=dict(metadata) if metadata is not None else None, + ) + return await asyncio.wrap_future(future) diff --git a/python/ebook_search/protected_phrases/store.py b/python/ebook_search/protected_phrases/store.py index 2a24f77..86aafd0 100644 --- a/python/ebook_search/protected_phrases/store.py +++ b/python/ebook_search/protected_phrases/store.py @@ -11,7 +11,11 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.dialects.sqlite import insert as sqlite_insert from python.ebook_search.protected_phrases.extraction import minimum_candidate_raw_count -from python.ebook_search.protected_phrases.models import PhraseCandidate, PhraseRecalculationResult +from python.ebook_search.protected_phrases.models import ( + CorpusPhraseStats, + PhraseCandidate, + PhraseRecalculationResult, +) from python.ebook_search.protected_phrases.text_normalization import normalize_text from python.orm.richie import ( EbookCandidatePhrase, @@ -27,7 +31,7 @@ if TYPE_CHECKING: from sqlalchemy.dialects.postgresql.dml import Insert as PostgresInsert from sqlalchemy.dialects.sqlite.dml import Insert as SqliteInsert - from sqlalchemy.orm import Session + from sqlalchemy.ext.asyncio import AsyncSession from python.ebook_search.config import EbookSearchConfig from python.ebook_search.protected_phrases.models import LLMJudgment @@ -36,14 +40,14 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def dialect_insert(session: Session, table: type[TableBase]) -> PostgresInsert | SqliteInsert: +def dialect_insert(session: AsyncSession, table: type[TableBase]) -> PostgresInsert | SqliteInsert: """Return a dialect-specific INSERT construct that supports ``ON CONFLICT DO UPDATE``. Production runs on PostgreSQL while tests run on SQLite; both support upserts with compatible SQLAlchemy constructs, so the correct one is chosen from the bound dialect. Args: - session (Session): Active database session whose bind selects the dialect. + session (AsyncSession): Active database session whose bind selects the dialect. table (type[TableBase]): Mapped table to insert into. Returns: @@ -54,35 +58,33 @@ def dialect_insert(session: Session, table: type[TableBase]) -> PostgresInsert | return pg_insert(table) - - -def load_book_text(session: Session, book_id: int) -> str: +async def load_book_text(session: AsyncSession, book_id: int) -> str: """Load a book's indexed chunk text as one string for phrase extraction. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. book_id (int): Book whose chunk text is loaded. Returns: str: The book's chunk text joined into a single string. """ - texts = session.scalars( + texts = await session.scalars( select(EbookChunk.text).where(EbookChunk.source_id == book_id).order_by(EbookChunk.chunk_index) ) return "\n\n".join(stripped for text in texts if (stripped := text.strip())) -def load_book_chapter_texts(session: Session, book_id: int) -> list[str]: +async def load_book_chapter_texts(session: AsyncSession, book_id: int) -> list[str]: """Reconstruct chapter-like text blocks from indexed chunks for phrase extraction. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. book_id (int): Book whose chunks are grouped into chapters. Returns: list[str]: Non-empty chapter-like text blocks in chunk order. """ - rows = session.execute( + rows = await session.execute( select(EbookChunk.chapter_id, EbookChunk.text) .where(EbookChunk.source_id == book_id) .order_by(EbookChunk.chunk_index) @@ -127,11 +129,11 @@ def metadata_for_source(source: EbookSource) -> dict[str, object | None]: } -def metadata_for_source_id(session: Session, source_id: int) -> dict[str, object | None]: +async def metadata_for_source_id(session: AsyncSession, source_id: int) -> dict[str, object | None]: """Return phrase extraction metadata for one indexed source by id. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. source_id (int): Id of the indexed source to read metadata from. Returns: @@ -140,70 +142,128 @@ def metadata_for_source_id(session: Session, source_id: int) -> dict[str, object Raises: ValueError: If no source exists with the given id. """ - source = session.get(EbookSource, source_id) + source = await session.get(EbookSource, source_id) if source is None: msg = f"No indexed source with id {source_id}" raise ValueError(msg) return metadata_for_source(source) -def count_protected_phrases(session: Session, book_id: int) -> int: +async def count_protected_phrases(session: AsyncSession, book_id: int) -> int: """Count stored protected phrases for one book. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. book_id (int): Book whose protected phrases are counted. Returns: int: Number of protected phrases stored for the book. """ - return session.scalars( - select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id) + return ( + await session.scalars( + select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id) + ) ).one() -def count_unjudged_candidates(session: Session, book_id: int, config: EbookSearchConfig) -> int: +async def count_unjudged_candidates(session: AsyncSession, book_id: int, config: EbookSearchConfig) -> int: """Count storable candidate rows for a book that have not yet been judged. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. book_id (int): Book whose unjudged candidates are counted. config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds. Returns: int: Number of storable, unjudged candidate rows for the book. """ - return session.scalars( - select(func.count(EbookCandidatePhrase.id)).where( - EbookCandidatePhrase.book_id == book_id, - EbookCandidatePhrase.llm_judged.is_(False), - EbookCandidatePhrase.token_count >= config.phrase_min_tokens, - EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config), + return ( + await session.scalars( + select(func.count(EbookCandidatePhrase.id)).where( + EbookCandidatePhrase.book_id == book_id, + EbookCandidatePhrase.llm_judged.is_(False), + EbookCandidatePhrase.token_count >= config.phrase_min_tokens, + EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config), + ) ) ).one() -def load_candidates_for_judgment( - session: Session, - book_id: int, - judgment_limit: int, - config: EbookSearchConfig, -) -> Sequence[EbookCandidatePhrase]: - """Load the top unjudged candidate rows for a book. - - Common-word phrases are filtered before storage (``filter_storable_candidates``), so - no re-check is needed here. +async def corpus_phrase_stats(session: AsyncSession) -> CorpusPhraseStats: + """Summarize candidate and protected phrase coverage across the whole corpus. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. + + Returns: + CorpusPhraseStats: Corpus-wide phrase counts and per-book coverage counts. + """ + total_books = (await session.scalars(select(func.count(EbookSource.id)))).one() + candidate_phrases, judged_candidates, books_with_candidates, books_with_unjudged = ( + await session.execute( + select( + func.count(EbookCandidatePhrase.id), + func.count(EbookCandidatePhrase.id).filter(EbookCandidatePhrase.llm_judged.is_(True)), + func.count(func.distinct(EbookCandidatePhrase.book_id)), + func.count(func.distinct(EbookCandidatePhrase.book_id)).filter( + EbookCandidatePhrase.llm_judged.is_(False) + ), + ) + ) + ).one() + protected_phrases = (await session.scalars(select(func.count(EbookProtectedPhrase.id)))).one() + return CorpusPhraseStats( + total_books=total_books, + books_with_candidates=books_with_candidates, + books_fully_judged=books_with_candidates - books_with_unjudged, + candidate_phrases=candidate_phrases, + judged_candidates=judged_candidates, + unjudged_candidates=candidate_phrases - judged_candidates, + protected_phrases=protected_phrases, + ) + + +async def book_ids_pending_first_judgment(session: AsyncSession) -> list[int]: + """Return books that have candidate phrases but no judged candidates yet. + + Args: + session (AsyncSession): Active database session. + + Returns: + list[int]: Book ids with candidates where judging has never run, ordered by id. + """ + judged_books = select(EbookCandidatePhrase.book_id).where(EbookCandidatePhrase.llm_judged.is_(True)).distinct() + return list( + ( + await session.scalars( + select(EbookCandidatePhrase.book_id) + .where(EbookCandidatePhrase.book_id.not_in(judged_books)) + .distinct() + .order_by(EbookCandidatePhrase.book_id) + ) + ).all() + ) + + +async def load_candidates_for_judgment( + session: AsyncSession, + book_id: int, + config: EbookSearchConfig, +) -> Sequence[EbookCandidatePhrase]: + """Load every storable unjudged candidate row for a book. + + Rows may have been stored before the current junk filters and score weights existed, so + callers re-check :func:`is_junk_phrase` and rescore before selecting what to judge. + + Args: + session (AsyncSession): Active database session. book_id (int): Book whose candidates are loaded. - judgment_limit (int): Maximum number of candidate rows to return. config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds. Returns: - Sequence[EbookCandidatePhrase]: Top storable, unjudged candidate rows ordered by score. + Sequence[EbookCandidatePhrase]: Storable, unjudged candidate rows ordered by stored score. """ - return session.scalars( + query = ( select(EbookCandidatePhrase) .where( EbookCandidatePhrase.book_id == book_id, @@ -216,8 +276,8 @@ def load_candidates_for_judgment( EbookCandidatePhrase.raw_count.desc(), EbookCandidatePhrase.id, ) - .limit(judgment_limit) - ).all() + ) + return (await session.scalars(query)).all() def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate: @@ -248,25 +308,23 @@ def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate: ) -def save_candidate_to_db( - session: Session, +def candidate_row_values( book_id: int, series_id: int | None, candidate: PhraseCandidate, *, judgment: LLMJudgment | None, -) -> EbookCandidatePhrase: - """Insert or update one candidate phrase row. +) -> dict[str, object]: + """Build the column values for one candidate phrase upsert. Args: - session (Session): Active database session. book_id (int): Book the candidate belongs to. series_id (int | None): Series scope stored on the row. candidate (PhraseCandidate): Candidate whose fields are written to the row. judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged. Returns: - EbookCandidatePhrase: The inserted or updated candidate row. + dict[str, object]: Column values keyed by column name. """ values: dict[str, object] = { "book_id": book_id, @@ -296,6 +354,30 @@ def save_candidate_to_db( llm_category=judgment.category, llm_reason=judgment.reason, ) + return values + + +async def save_candidate_to_db( + session: AsyncSession, + book_id: int, + series_id: int | None, + candidate: PhraseCandidate, + *, + judgment: LLMJudgment | None, +) -> EbookCandidatePhrase: + """Insert or update one candidate phrase row. + + Args: + session (AsyncSession): Active database session. + book_id (int): Book the candidate belongs to. + series_id (int | None): Series scope stored on the row. + candidate (PhraseCandidate): Candidate whose fields are written to the row. + judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged. + + Returns: + EbookCandidatePhrase: The inserted or updated candidate row. + """ + values = candidate_row_values(book_id, series_id, candidate, judgment=judgment) # Preserve an existing judgment when this call is only refreshing candidate fields. skip_update = {"book_id", "phrase_norm"} @@ -306,11 +388,94 @@ def save_candidate_to_db( index_elements=["book_id", "phrase_norm"], set_={column: insert_statement.excluded[column] for column in values if column not in skip_update}, ).returning(EbookCandidatePhrase) - return session.scalars(statement, execution_options={"populate_existing": True}).one() + return (await session.scalars(statement, execution_options={"populate_existing": True})).one() -def upsert_protected_phrase( - session: Session, +BULK_CANDIDATE_UPSERT_CHUNK = 1000 + + +async def bulk_upsert_unjudged_candidates( + session: AsyncSession, + book_id: int, + series_id: int | None, + candidates: Sequence[PhraseCandidate], +) -> int: + """Insert or update many freshly extracted candidate rows in chunked multi-row upserts. + + Saving one row per statement costs one database round trip per candidate, which dominated + generation time for full books, so candidates are written ``BULK_CANDIDATE_UPSERT_CHUNK`` + rows per statement instead. Existing judgments and sample contexts are never overwritten: + fresh extractions carry no contexts, and ``llm_judged`` plus the ``llm_*`` columns are left + out of the conflict update. Candidates must have unique ``phrase_norm`` values, as produced + by extraction, since one multi-row upsert cannot touch the same row twice. + + Args: + session (Session): Active database session. + book_id (int): Book the candidates belong to. + series_id (int | None): Series scope stored on the rows. + candidates (Sequence[PhraseCandidate]): Freshly extracted candidates to persist. + + Returns: + int: Number of candidate rows written. + """ + values = [ + candidate_row_values(book_id, series_id, candidate, judgment=None) + for candidate in candidates + if not candidate.sample_contexts + ] + if len(values) != len(candidates): + msg = "bulk_upsert_unjudged_candidates only accepts freshly extracted candidates without sample contexts" + raise ValueError(msg) + skip_update = {"book_id", "phrase_norm", "llm_judged"} + for chunk_start in range(0, len(values), BULK_CANDIDATE_UPSERT_CHUNK): + chunk = values[chunk_start : chunk_start + BULK_CANDIDATE_UPSERT_CHUNK] + insert_statement = dialect_insert(session, EbookCandidatePhrase).values(chunk) + statement = insert_statement.on_conflict_do_update( + index_elements=["book_id", "phrase_norm"], + set_={column: insert_statement.excluded[column] for column in chunk[0] if column not in skip_update}, + ) + await session.execute(statement) + return len(values) + + +def new_candidate_row(book_id: int, series_id: int | None, candidate: PhraseCandidate) -> EbookCandidatePhrase: + """Build a fresh unjudged candidate row without checking for an existing one. + + Unlike :func:`save_candidate_to_db`, this does no lookup, so it is only safe when the caller + guarantees there is no existing row for ``(book_id, candidate.phrase_norm)`` — for example + right after :func:`delete_phrase_data_for_book` has cleared the book. + + Args: + book_id (int): Book the candidate belongs to. + series_id (int | None): Series scope stored on the row. + candidate (PhraseCandidate): Candidate whose fields are written to the row. + + Returns: + EbookCandidatePhrase: A new, unattached candidate row. + """ + row = EbookCandidatePhrase(book_id=book_id, phrase_norm=candidate.phrase_norm) + row.llm_judged = False + row.series_id = series_id + row.phrase_text = candidate.phrase_text + row.token_count = candidate.token_count + row.source_raw_ngram = candidate.source_raw_ngram + row.source_yake = candidate.source_yake + row.source_spacy_ner = candidate.source_spacy_ner + row.source_spacy_noun_chunk = candidate.source_spacy_noun_chunk + row.source_capitalized = candidate.source_capitalized + row.source_metadata = candidate.source_metadata + row.spacy_label = candidate.spacy_label + row.raw_count = candidate.raw_count + row.chapter_count = candidate.chapter_count + row.yake_score = candidate.yake_score + row.candidate_score = candidate.candidate_score + if candidate.sample_contexts: + row.sample_contexts = list(candidate.sample_contexts) + return row + + +async def upsert_protected_phrase( + session: AsyncSession, book_id: int, series_id: int | None, candidate: PhraseCandidate, @@ -320,7 +485,7 @@ def upsert_protected_phrase( """Insert or update one accepted protected phrase and its aliases. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. book_id (int): Book the protected phrase belongs to. series_id (int | None): Series scope stored on the phrase. candidate (PhraseCandidate): Candidate the phrase was promoted from. @@ -360,18 +525,22 @@ def upsert_protected_phrase( column: insert_statement.excluded[column] for column in values if column not in {"book_id", "phrase_norm"} }, ).returning(EbookProtectedPhrase) - row = session.scalars(statement, execution_options={"populate_existing": True}).one() + row = (await session.scalars(statement, execution_options={"populate_existing": True})).one() for alias_text in judgment.aliases: - upsert_phrase_alias(session, row, alias_text) + await upsert_phrase_alias(session, row, alias_text) return row -def upsert_phrase_alias(session: Session, phrase: EbookProtectedPhrase, alias_text: str) -> EbookPhraseAlias | None: +async def upsert_phrase_alias( + session: AsyncSession, + phrase: EbookProtectedPhrase, + alias_text: str, +) -> EbookPhraseAlias | None: """Insert or update one protected phrase alias. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. phrase (EbookProtectedPhrase): Protected phrase the alias points to. alias_text (str): Alias surface form to store. @@ -395,7 +564,7 @@ def upsert_phrase_alias(session: Session, phrase: EbookProtectedPhrase, alias_te "confidence": insert_statement.excluded.confidence, }, ).returning(EbookPhraseAlias) - return session.scalars(statement, execution_options={"populate_existing": True}).one() + return (await session.scalars(statement, execution_options={"populate_existing": True})).one() def make_canonical_id(judgment: LLMJudgment, phrase_norm: str) -> str: @@ -426,15 +595,15 @@ def slugify_identifier(value: str) -> str: return slug.strip("_") or "unknown" -def prune_unstorable_unjudged_candidate_phrases( - session: Session, +async def prune_unstorable_unjudged_candidate_phrases( + session: AsyncSession, book_id: int, config: EbookSearchConfig, ) -> int: """Delete old unjudged candidate rows that no longer satisfy storage filters. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. book_id (int): Book whose stale candidates are pruned. config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds. @@ -442,7 +611,7 @@ def prune_unstorable_unjudged_candidate_phrases( int: Number of candidate rows deleted. """ deleted = rowcount( - session.execute( + await session.execute( delete(EbookCandidatePhrase).where( EbookCandidatePhrase.book_id == book_id, EbookCandidatePhrase.llm_judged.is_(False), @@ -464,40 +633,42 @@ def prune_unstorable_unjudged_candidate_phrases( return deleted -def delete_phrase_data_for_book(session: Session, book_id: int) -> PhraseRecalculationResult: +async def delete_phrase_data_for_book(session: AsyncSession, book_id: int) -> PhraseRecalculationResult: """Delete all candidate, protected, alias, and mention phrase data for one book. Args: - session (Session): Active database session. + session (AsyncSession): Active database session. book_id (int): Book whose phrase data is deleted. Returns: PhraseRecalculationResult: Deleted-row counts with ``candidate_phrases`` set to 0. """ - protected_ids = session.scalars( - select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id) + protected_ids = ( + await session.scalars(select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id)) ).all() deleted_aliases = 0 if protected_ids: deleted_aliases = rowcount( - session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids))) + await session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids))) ) deleted_mentions = rowcount( - session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id)) + await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id)) ) if protected_ids: deleted_mentions += rowcount( - session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids))) + await session.execute( + delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids)) + ) ) deleted_protected = rowcount( - session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id)) + await session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id)) ) deleted_candidates = rowcount( - session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id)) + await session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id)) ) - session.flush() + await session.flush() logger.info( "ebook_candidate_phrase_data_deleted book_id=%s candidates=%s protected=%s aliases=%s mentions=%s", book_id, @@ -515,6 +686,7 @@ def delete_phrase_data_for_book(session: Session, book_id: int) -> PhraseRecalcu candidate_phrases=0, ) + def rowcount(result: object) -> int: """Return a safe integer rowcount from a SQLAlchemy execution result. diff --git a/python/ebook_search/rerank.py b/python/ebook_search/rerank.py index 05354be..b47e5ec 100644 --- a/python/ebook_search/rerank.py +++ b/python/ebook_search/rerank.py @@ -9,6 +9,8 @@ from typing import TYPE_CHECKING from python.ebook_search.llm_interface import request_rerank if TYPE_CHECKING: + import httpx + from python.ebook_search.config import RerankConfig from python.ebook_search.search import SearchResult @@ -23,7 +25,12 @@ class RerankResult: score: float -def rerank_chunks(query: str, candidates: list[SearchResult], config: RerankConfig) -> list[SearchResult]: +async def rerank_chunks( + client: httpx.AsyncClient, + query: str, + candidates: list[SearchResult], + config: RerankConfig, +) -> list[SearchResult]: """Rerank candidates with a vLLM rerank endpoint.""" if not candidates: return [] @@ -34,7 +41,7 @@ def rerank_chunks(query: str, candidates: list[SearchResult], config: RerankConf config.model, len(candidates), ) - scores = score_candidates(query, candidates, config) + scores = await score_candidates(client, query, candidates, config) results = sorted( ( replace( @@ -56,13 +63,14 @@ def rerank_chunks(query: str, candidates: list[SearchResult], config: RerankConf return results -def score_candidates( +async def score_candidates( + client: httpx.AsyncClient, query: str, candidates: list[SearchResult], config: RerankConfig, ) -> dict[int, RerankResult]: """Score candidate chunks with the configured rerank API.""" - body = request_rerank(query, [candidate.text for candidate in candidates], config) + body = await request_rerank(client, query, [candidate.text for candidate in candidates], config) if body is None: return zero_rerank_scores(candidates) diff --git a/python/ebook_search/search.py b/python/ebook_search/search.py index c59f87a..d8a5709 100644 --- a/python/ebook_search/search.py +++ b/python/ebook_search/search.py @@ -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", diff --git a/python/ebook_search/timing.py b/python/ebook_search/timing.py index eb8e474..9dff9a6 100644 --- a/python/ebook_search/timing.py +++ b/python/ebook_search/timing.py @@ -7,7 +7,7 @@ from time import perf_counter from typing import TYPE_CHECKING if TYPE_CHECKING: - from collections.abc import Callable + from collections.abc import Awaitable, Callable @dataclass(frozen=True) @@ -34,3 +34,10 @@ def timed_result[T, **P]( start_seconds = perf_counter() result = operation(*args, **kwargs) return result, runtime_step_from_start(name, start_seconds) + + +async def async_timed_result[T](name: str, awaitable: Awaitable[T]) -> tuple[T, RuntimeStep]: + """Await an operation and return its result plus elapsed runtime.""" + start_seconds = perf_counter() + result = await awaitable + return result, runtime_step_from_start(name, start_seconds) diff --git a/python/fastapi_tools/__init__.py b/python/fastapi_tools/__init__.py index d55eb3d..8490ccb 100644 --- a/python/fastapi_tools/__init__.py +++ b/python/fastapi_tools/__init__.py @@ -1,6 +1,6 @@ """Reusable FastAPI tools.""" -from python.fastapi_tools.db import DbSession, get_db +from python.fastapi_tools.db import AsyncDbSession, DbSession, get_async_db, get_db from python.fastapi_tools.zstd_middleware import ZstdMiddleware -__all__ = ["DbSession", "ZstdMiddleware", "get_db"] +__all__ = ["AsyncDbSession", "DbSession", "ZstdMiddleware", "get_async_db", "get_db"] diff --git a/python/fastapi_tools/db.py b/python/fastapi_tools/db.py index c6ee685..6dad25a 100644 --- a/python/fastapi_tools/db.py +++ b/python/fastapi_tools/db.py @@ -5,10 +5,11 @@ from __future__ import annotations from typing import TYPE_CHECKING, Annotated from fastapi import Depends, Request +from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import AsyncIterator, Iterator def get_db(request: Request) -> Iterator[Session]: @@ -17,4 +18,15 @@ def get_db(request: Request) -> Iterator[Session]: yield session +async def get_async_db(request: Request) -> AsyncIterator[AsyncSession]: + """Get an async database session from app state. + + expire_on_commit=False keeps ORM attributes readable after commit without + triggering implicit IO, which would raise under asyncio. + """ + async with AsyncSession(request.app.state.engine, expire_on_commit=False) as session: + yield session + + DbSession = Annotated[Session, Depends(get_db)] +AsyncDbSession = Annotated[AsyncSession, Depends(get_async_db)] diff --git a/python/orm/common.py b/python/orm/common.py index 2b42564..fcf7da4 100644 --- a/python/orm/common.py +++ b/python/orm/common.py @@ -7,6 +7,7 @@ from typing import cast from sqlalchemy import create_engine from sqlalchemy.engine import URL, Engine +from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine NAMING_CONVENTION = { "ix": "ix_%(table_name)s_%(column_0_name)s", @@ -31,6 +32,36 @@ def get_connection_info(name: str) -> tuple[str, str, str, str, str | None]: return cast("tuple[str, str, str, str, str | None]", (database, host, port, username, password)) +def build_postgres_url(name: str, *, vector_engine: bool = False) -> tuple[URL, dict[str, str]]: + """Build the Postgres connection URL and connect_args from environment variables. + + Args: + name (str): The name of the environment variable prefix. + vector_engine (bool, optional): Whether to use the vector search schema. Defaults to False. + This updates the search path to include the vector types and operators. + + Returns: + tuple[URL, dict[str, str]]: The SQLAlchemy URL and connect_args for create_engine. + """ + database, host, port, username, password = get_connection_info(name) + + url = URL.create( + drivername="postgresql+psycopg", + username=username, + password=password, + host=host, + port=int(port), + database=database, + ) + + connect_args = {} + # There more better way to do this is with separate PG account and a dedicated vector schema for the vector types + if vector_engine: + connect_args["options"] = "-csearch_path=main,public" + + return url, connect_args + + def get_postgres_engine( *, name: str = "POSTGRES", @@ -51,21 +82,7 @@ def get_postgres_engine( Returns: Engine: The SQLAlchemy engine. """ - database, host, port, username, password = get_connection_info(name) - - url = URL.create( - drivername="postgresql+psycopg", - username=username, - password=password, - host=host, - port=int(port), - database=database, - ) - - connect_args = {} - # There more better way to do this is with separate PG account and a dedicated vector schema for the vector types - if vector_engine: - connect_args["options"] = "-csearch_path=main,public" + url, connect_args = build_postgres_url(name, vector_engine=vector_engine) return create_engine( url=url, @@ -74,3 +91,34 @@ def get_postgres_engine( connect_args=connect_args, pool_size=pool_size, ) + + +def get_async_postgres_engine( + *, + name: str = "POSTGRES", + pool_pre_ping: bool = True, + vector_engine: bool = False, + pool_size: int = 8, +) -> AsyncEngine: + """Create an async SQLAlchemy engine from environment variables. + + Args: + name (str, optional): The name of the environment variable prefix. Defaults to "POSTGRES". + pool_pre_ping (bool, optional): Whether to ping the database before each connection. Defaults to True. + This fixes the issue of trying to use a conection that has timed out on the database side. + vector_engine (bool, optional): Whether to use the vector search schema. Defaults to False. + This updates the search path the incldued the vecore types and operators. + pool_size (int, optional): Number of connections to keep in the pool. Defaults to 8. + + Returns: + AsyncEngine: The async SQLAlchemy engine. + """ + url, connect_args = build_postgres_url(name, vector_engine=vector_engine) + + return create_async_engine( + url=url, + pool_pre_ping=pool_pre_ping, + pool_recycle=1800, + connect_args=connect_args, + pool_size=pool_size, + )