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
+
+
+
+
+ | Candidates |
+ Judged |
+ Unjudged |
+ Protected |
+ Books indexed |
+ Books generated |
+ Books 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,
+ )