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

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

Browse Source 
Convert the ebook-search web app to async end to end and add concurrency
to the protected-phrase extraction and judging pipeline so large books no
longer block the event loop or the UI.

ORM / infra:
- Add get_async_postgres_engine and factor shared URL/connect_args building
  into build_postgres_url (reused by the sync and async engine builders)
- Add async FastAPI session helpers (get_async_db, AsyncDbSession) with
  expire_on_commit=False to avoid implicit IO under asyncio

App:
- Use AsyncEngine/AsyncSession throughout routes, search, ingest, embeddings,
  answer, rerank and LLM calls; convert handlers to async
- Share a single httpx.AsyncClient in app state for LLM requests; size the
  connection pool for concurrent phrase-judging workers
- Add judge_tasks: run per-book judging as tracked background tasks so a
  book already being judged isn't double-queued

Protected phrases:
- Add a process pool (pool.py) and worker-count config
  (extraction/judge book/phrase workers) to parallelize candidate generation
  and judging
- Split admin actions into all/missing variants for generation and judging

Config:
- Add protected_phrase_extraction_workers, phrase_judge_book_workers,
  phrase_judge_phrase_workers
This commit is contained in:
Richie
2026-07-24 11:38:50 -04:00
parent 38c01ec121
commit 2706c4417d
29 changed files with 1824 additions and 769 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ebook_search/answer.py
+10 -2
View File
@@ -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,
[
{
python/ebook_search/api/bm25_tasks.py
+32 -19
View File
@@ -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")
python/ebook_search/api/dependencies.py
+10 -3
View File
@@ -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)]
python/ebook_search/api/judge_tasks.py
+131
View File
@@ -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)
python/ebook_search/api/main.py
+17 -6
View File
@@ -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:
python/ebook_search/api/routes/admin.py
+127 -33
View File
@@ -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(
python/ebook_search/api/routes/health.py
+15 -13
View File
@@ -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
python/ebook_search/api/routes/page.py
+69 -28
View File
@@ -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)
python/ebook_search/api/routes/search.py
+11 -5
View File
@@ -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,
python/ebook_search/api/templates/admin.html
+73 -28
View File
@@ -8,37 +8,25 @@ head %}
<form hx-post="/admin/scan" hx-target="#admin-status" hx-swap="innerHTML">
<button type="submit">Scan</button>
</form>
<form
hx-post="/admin/generate-ngrams"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Generate n-grams</button>
</form>
<form
hx-post="/admin/judge-ngrams"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Judge n-grams</button>
</form>
<form
hx-post="/admin/embed-missing"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Embed</button>
</form>
<form
hx-post="/admin/embed-all"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Embed all</button>
</form>
</section>
<section>
<h2>Embeddings</h2>
<section class="actions">
<form
hx-post="/admin/embed-missing"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Embed</button>
</form>
<form
hx-post="/admin/embed-all"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Embed all</button>
</form>
</section>
<table>
<thead>
<tr>
@@ -62,4 +50,61 @@ head %}
</tbody>
</table>
</section>
<section>
<h2>Protected phrases</h2>
<section class="actions">
<form
hx-post="/admin/phrases/generate-all"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Regenerate all phrases</button>
</form>
<form
hx-post="/admin/phrases/generate-missing"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Add missing phrases</button>
</form>
<form
hx-post="/admin/phrases/judge-all"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Judge all phrases</button>
</form>
<form
hx-post="/admin/phrases/judge-missing"
hx-target="#admin-status"
hx-swap="innerHTML"
>
<button type="submit">Judge missing phrases</button>
</form>
</section>
<table>
<thead>
<tr>
<th>Candidates</th>
<th>Judged</th>
<th>Unjudged</th>
<th>Protected</th>
<th>Books indexed</th>
<th>Books generated</th>
<th>Books fully judged</th>
</tr>
</thead>
<tbody>
<tr>
<td>{{ phrase_stats.candidate_phrases }}</td>
<td>{{ phrase_stats.judged_candidates }}</td>
<td>{{ phrase_stats.unjudged_candidates }}</td>
<td>{{ phrase_stats.protected_phrases }}</td>
<td>{{ phrase_stats.total_books }}</td>
<td>{{ phrase_stats.books_with_candidates }}</td>
<td>{{ phrase_stats.books_fully_judged }}</td>
</tr>
</tbody>
</table>
</section>
{% endblock %}
python/ebook_search/api/templates/book_detail.html
+9
View File
@@ -30,6 +30,15 @@
>
<button type="submit">Recalculate phrases</button>
</form>
<form
method="post"
action="/books/{{ source.id }}/judge-phrases"
onsubmit="return confirm('Judge candidate phrases for this book with the LLM?');"
>
<button type="submit"{% if judging_in_progress %} disabled{% endif %}>
{% if judging_in_progress %}Judging&hellip;{% else %}Judge phrases{% endif %}
</button>
</form>
<section>
<h2>Candidate n-grams</h2>
python/ebook_search/bm25_corpus.py
+19 -15
View File
@@ -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(
python/ebook_search/config.py
+3
View File
@@ -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
python/ebook_search/embeddings.py
+24 -19
View File
@@ -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)
python/ebook_search/ingest.py
+43 -24
View File
@@ -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,
python/ebook_search/llm_interface.py
+63 -13
View File
@@ -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={
python/ebook_search/protected_phrases/extraction.py
+185 -79
View File
@@ -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),
python/ebook_search/protected_phrases/generate_ngrams.py
+191 -95
View File
@@ -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
python/ebook_search/protected_phrases/judge_ngrams.py
+235 -205
View File
@@ -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:
python/ebook_search/protected_phrases/matching.py
+37 -33
View File
@@ -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()}
python/ebook_search/protected_phrases/models.py
+23
View File
@@ -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.
python/ebook_search/protected_phrases/pool.py
+101
View File
@@ -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)
python/ebook_search/protected_phrases/store.py
+244 -72
View File
@@ -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.
python/ebook_search/rerank.py
+12 -4
View File
@@ -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)
python/ebook_search/search.py
+54 -54
View File
@@ -2,17 +2,17 @@
from __future__ import annotations
import asyncio
import logging
import re
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, replace
from typing import TYPE_CHECKING
from pgvector.sqlalchemy import Vector
from sqlalchemy import literal, select
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import Session
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.bm25_corpus import (
BM25CorpusUnavailableError,
@@ -25,7 +25,7 @@ from python.ebook_search.protected_phrases.matching import (
phrase_hits_for_chunks,
)
from python.ebook_search.rerank import rerank_chunks
from python.ebook_search.timing import RuntimeStep, timed_result
from python.ebook_search.timing import RuntimeStep, async_timed_result, timed_result
from python.orm.richie import (
EbookChapter,
EbookChunk,
@@ -36,7 +36,8 @@ from python.orm.richie import (
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from sqlalchemy.engine import Engine
import httpx
from sqlalchemy.ext.asyncio import AsyncEngine
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import HydratedPhraseMatch
@@ -90,8 +91,9 @@ class RetrievalResponse:
timings: tuple[RuntimeStep, ...]
def search_ebooks(
engine: Engine,
async def search_ebooks(
engine: AsyncEngine,
client: httpx.AsyncClient,
query: str,
config: EbookSearchConfig,
*,
@@ -112,16 +114,15 @@ def search_ebooks(
)
timings: list[RuntimeStep] = []
if phrase_matching_enabled:
phrase_matches, timing = timed_result("Protected phrase detection", query_phrase_matches, engine, query, config)
phrase_matches, timing = await async_timed_result(
"Protected phrase detection", query_phrase_matches(engine, query, config)
)
else:
phrase_matches, timing = timed_result("Protected phrase detection skipped", skip_phrase_matches)
timings.append(timing)
retrieval, timing = timed_result(
retrieval, timing = await async_timed_result(
"Hybrid retrieval",
parallel_retrieval,
engine,
query,
config,
parallel_retrieval(engine, client, query, config),
)
timings.extend(retrieval.timings)
timings.append(timing)
@@ -134,19 +135,15 @@ def search_ebooks(
)
timings.append(timing)
if phrase_matching_enabled:
fused, timing = timed_result(
fused, timing = await async_timed_result(
"Phrase mention boost",
apply_phrase_mention_boosts,
engine,
fused,
phrase_matches,
config.phrase_hit_boost,
apply_phrase_mention_boosts(engine, fused, phrase_matches, config.phrase_hit_boost),
)
else:
fused, timing = timed_result("Phrase mention boost skipped", skip_phrase_mention_boosts, fused)
timings.append(timing)
if config.rerank.enabled and rerank:
response, timing = timed_result("Rerank", apply_rerank, query, fused, config)
response, timing = await async_timed_result("Rerank", apply_rerank(client, query, fused, config))
else:
response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
timings.append(timing)
@@ -172,11 +169,15 @@ def skip_phrase_matches() -> list[HydratedPhraseMatch]:
return []
def query_phrase_matches(engine: Engine, query: str, config: EbookSearchConfig) -> list[HydratedPhraseMatch]:
async def query_phrase_matches(
engine: AsyncEngine,
query: str,
config: EbookSearchConfig,
) -> list[HydratedPhraseMatch]:
"""Detect protected phrases in a query without making search fail when phrase tables are unavailable."""
try:
with Session(engine) as session:
return detect_protected_phrases_for_query(session, query, config)
async with AsyncSession(engine) as session:
return await detect_protected_phrases_for_query(session, query, config)
except SQLAlchemyError as error:
logger.warning("ebook_protected_phrase_detection_unavailable error=%s", error)
return []
@@ -188,8 +189,8 @@ def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchRes
return candidates
def apply_phrase_mention_boosts(
engine: Engine,
async def apply_phrase_mention_boosts(
engine: AsyncEngine,
candidates: list[SearchResult],
phrase_matches: Sequence[HydratedPhraseMatch],
phrase_hit_boost: float,
@@ -201,8 +202,8 @@ def apply_phrase_mention_boosts(
chunk_ids = [candidate.chunk_id for candidate in candidates]
try:
with Session(engine) as session:
phrase_hits = phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
async with AsyncSession(engine) as session:
phrase_hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
except SQLAlchemyError as error:
logger.warning("ebook_phrase_boost_unavailable error=%s", error)
return candidates
@@ -241,30 +242,21 @@ def phrase_rank_source(rank_source: str, phrase_hit_count: int) -> str:
return f"{rank_source} + phrases"
def parallel_retrieval(
engine: Engine,
async def parallel_retrieval(
engine: AsyncEngine,
client: httpx.AsyncClient,
query: str,
config: EbookSearchConfig,
) -> RetrievalResponse:
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions."""
with ThreadPoolExecutor(max_workers=2, thread_name_prefix="ebook-search") as executor:
vector_future = executor.submit(
timed_result,
"Embedding + vector search",
vector_candidates,
engine,
query,
config,
)
bm25_future = executor.submit(
timed_result,
"BM25 search",
bm25_candidates,
query,
config,
)
vector_results, vector_timing = vector_future.result()
lexical_results, lexical_timing = bm25_future.result()
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions.
BM25 scoring is pure CPU work over the cached corpus, so it runs in a worker thread
instead of on the event loop.
"""
(vector_results, vector_timing), (lexical_results, lexical_timing) = await asyncio.gather(
async_timed_result("Embedding + vector search", vector_candidates(engine, client, query, config)),
async_timed_result("BM25 search", asyncio.to_thread(bm25_candidates, query, config)),
)
logger.info(
"ebook_parallel_retrieval_complete vector_candidates=%s lexical_candidates=%s",
@@ -291,13 +283,14 @@ def skip_rerank(
return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
def apply_rerank(
async def apply_rerank(
client: httpx.AsyncClient,
query: str,
candidates: list[SearchResult],
config: EbookSearchConfig,
) -> SearchResponse:
"""Rerank already-fused hybrid candidates."""
reranked = rerank_chunks(query, candidates[: config.rerank.candidates], config.rerank)
reranked = await rerank_chunks(client, query, candidates[: config.rerank.candidates], config.rerank)
logger.info(
"ebook_rerank_complete input_candidates=%s returned=%s",
min(len(candidates), config.rerank.candidates),
@@ -310,10 +303,17 @@ def apply_rerank(
)
def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) -> list[SearchResult]:
async def vector_candidates(
engine: AsyncEngine,
client: httpx.AsyncClient,
query: str,
config: EbookSearchConfig,
) -> list[SearchResult]:
"""Return pgvector cosine candidates for a natural-language query."""
with Session(engine) as session:
model = session.scalar(select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model))
async with AsyncSession(engine) as session:
model = await session.scalar(
select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model)
)
if model is None:
msg = f"Embedding model is not registered: {config.embedding_model}"
raise ValueError(msg)
@@ -323,7 +323,7 @@ def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) ->
msg = f"Model row dimension {model.dimension} does not match configured dimension {expected_dimension}"
raise ValueError(msg)
embedding = embed_query(query, config)
embedding = await embed_query(client, query, config)
limit = max(config.rerank.candidates, config.top_k) * config.vector_candidate_multiplier
embedding_table = get_embedding_table(model.dimension)
@@ -349,7 +349,7 @@ def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) ->
.order_by(distance)
.limit(limit)
)
rows = session.execute(statement).mappings()
rows = (await session.execute(statement)).mappings()
results = [search_result_from_row(row) for row in rows]
logger.info(
"ebook_vector_search_complete model=%s dimension=%s candidates=%s",
python/ebook_search/timing.py
+8 -1
View File
@@ -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)
python/fastapi_tools/__init__.py
+2 -2
View File
@@ -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"]
python/fastapi_tools/db.py
+13 -1
View File
@@ -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)]
python/orm/common.py
+63 -15
View File
@@ -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,
)