Compare commits

..
Author SHA1 Message Date
Richie 659b1e2798 fix(protected-phrases): isolate phrase generation per book
treefmt / nix fmt (pull_request) Failing after 5s
pytest / pytest (pull_request) Failing after 31s
build_systems / build-brain (pull_request) Successful in 49s
build_systems / build-bob (pull_request) Successful in 49s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m1s
build_systems / build-jeeves (pull_request) Successful in 2m39s
Run full-book candidate generation inside worker-owned sessions so each book commits independently during backfills. Abort recalculation when a book has no indexed chapters to preserve existing phrase data, and update admin/UI tests for the new generation flow.
2026-07-11 21:10:32 -04:00
Richie b13bb39401 feat(admin): simplify phrase generation by removing missing phrases endpoint 2026-07-09 23:08:04 -04:00
Richie ffdb93d352 refactor(ebook): remove spaCy-ner attributes from PhraseCandidate and related functions 2026-07-09 23:05:08 -04:00
Richie 0cf05c52b9 ran treefmt 2026-07-09 22:09:29 -04:00
15 changed files with 359 additions and 449 deletions
+1 -3
View File
@@ -17,9 +17,7 @@
python-env = final: _prev: {
my_python = final.python314.withPackages (
ps:
with ps;
[
ps: with ps; [
alembic
apprise
apscheduler
@@ -0,0 +1,55 @@
"""remove spaCy-ner.
Revision ID: 751260fc3228
Revises: dddee09eddcc
Create Date: 2026-07-09 23:03:39.554083
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import sqlalchemy as sa
from alembic import op
from python.orm import RichieBase
if TYPE_CHECKING:
from collections.abc import Sequence
# revision identifiers, used by Alembic.
revision: str = "751260fc3228"
down_revision: str | None = "dddee09eddcc"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
schema = RichieBase.schema_name
def upgrade() -> None:
"""Upgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column("candidate_phrases", "source_spacy_noun_chunk", schema=schema)
op.drop_column("candidate_phrases", "source_spacy_ner", schema=schema)
op.drop_column("candidate_phrases", "spacy_label", schema=schema)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade."""
# ### commands auto generated by Alembic - please adjust! ###
op.add_column(
"candidate_phrases", sa.Column("spacy_label", sa.VARCHAR(), autoincrement=False, nullable=True), schema=schema
)
op.add_column(
"candidate_phrases",
sa.Column("source_spacy_ner", sa.BOOLEAN(), autoincrement=False, nullable=False),
schema=schema,
)
op.add_column(
"candidate_phrases",
sa.Column("source_spacy_noun_chunk", sa.BOOLEAN(), autoincrement=False, nullable=False),
schema=schema,
)
# ### end Alembic commands ###
+4 -40
View File
@@ -61,56 +61,20 @@ async def scan_library(request: Request, config: AppConfig, session: AsyncDbSess
@router.post("/phrases/generate-all", response_class=HTMLResponse)
async def generate_all_phrases(request: Request, config: AppConfig, session: AsyncDbSession) -> HTMLResponse:
async def generate_all_phrases(request: Request, config: AppConfig, engine: AppEngine) -> 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 = await generate_candidate_phrases_for_books(session, config, only_missing=only_missing)
await session.commit()
result = await generate_candidate_phrases_for_books(engine, config)
except Exception as error:
await session.rollback()
logger.exception("ebook_admin_generate_phrases_failed only_missing=%s", only_missing)
logger.exception("ebook_admin_generate_phrases_failed")
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
logger.info(
"ebook_admin_generate_phrases_complete only_missing=%s books_seen=%s books_built=%s candidates=%s",
only_missing,
"ebook_admin_generate_phrases_complete books_seen=%s books_built=%s candidates=%s",
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",
+4 -1
View File
@@ -168,7 +168,10 @@ async def recalculate_book_phrases(source_id: int, config: AppConfig, session: A
if source is None:
raise HTTPException(status_code=404, detail="Book not found")
result = await recalculate_candidate_phrases_for_book(session, source, config, use_process_pool=True)
try:
result = await recalculate_candidate_phrases_for_book(session, source, config)
except ValueError as error:
raise HTTPException(status_code=409, detail=str(error)) from error
logger.info(
"ebook_book_phrase_recalculation_complete source_id=%s candidates=%s deleted_candidates=%s "
"deleted_protected=%s deleted_aliases=%s deleted_mentions=%s",
@@ -60,13 +60,6 @@ head %}
>
<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"
@@ -36,32 +36,6 @@ 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}")
class SpacySpan(Protocol):
"""Small protocol for the spaCy span attributes used by this module."""
text: str
class SpacyEntity(SpacySpan, Protocol):
"""Small protocol for the spaCy entity attributes used by this module."""
label_: str
class SpacyDoc(Protocol):
"""Small protocol for the spaCy doc attributes used by this module."""
ents: Iterable[SpacyEntity]
noun_chunks: Iterable[SpacySpan]
class SpacyLanguage(Protocol):
"""Small protocol for a callable spaCy language pipeline."""
def __call__(self, text: str) -> SpacyDoc:
"""Parse text into a spaCy-like doc."""
class YakeExtractor(Protocol):
"""Small protocol for the YAKE extractor used by this module."""
@@ -242,55 +216,6 @@ def extract_yake_candidates(
return out
def extract_spacy_candidates(
book_text: str,
nlp: SpacyLanguage,
config: EbookSearchConfig,
) -> dict[str, PhraseCandidate]:
"""Extract spaCy named entities and noun chunks from one text block.
Args:
book_text (str): Text block to parse with spaCy.
nlp (SpacyLanguage): Callable spaCy language pipeline.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from entities and noun chunks.
"""
out: dict[str, PhraseCandidate] = {}
doc = nlp(book_text)
for ent in doc.ents:
normalized = normalize_candidate_phrase(
ent.text,
config,
max_tokens=config.phrase_max_entity_tokens,
)
if normalized is None:
continue
phrase_text, phrase_norm, token_count = normalized
out[phrase_norm] = PhraseCandidate(
phrase_text=phrase_text,
phrase_norm=phrase_norm,
token_count=token_count,
source_spacy_ner=True,
spacy_label=ent.label_,
)
for chunk in doc.noun_chunks:
normalized = normalize_candidate_phrase(chunk.text, config, strip_leading_article=True)
if normalized is None:
continue
phrase_text, phrase_norm, token_count = normalized
out[phrase_norm] = PhraseCandidate(
phrase_text=phrase_text,
phrase_norm=phrase_norm,
token_count=token_count,
source_spacy_noun_chunk=True,
)
return out
def extract_capitalized_phrases(original_text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]:
"""Extract capitalized phrase runs that often carry fictional terms.
@@ -392,16 +317,12 @@ def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None:
"""
existing.source_raw_ngram = existing.source_raw_ngram or item.source_raw_ngram
existing.source_yake = existing.source_yake or item.source_yake
existing.source_spacy_ner = existing.source_spacy_ner or item.source_spacy_ner
existing.source_spacy_noun_chunk = existing.source_spacy_noun_chunk or item.source_spacy_noun_chunk
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:
existing.spacy_label = item.spacy_label
def enrich_with_frequency_and_chapter_counts(
@@ -599,8 +520,6 @@ def non_raw_source_count(candidate: PhraseCandidate) -> int:
return sum(
(
candidate.source_yake,
candidate.source_spacy_ner,
candidate.source_spacy_noun_chunk,
candidate.source_capitalized,
candidate.source_metadata,
)
@@ -646,8 +565,6 @@ def source_score(candidate: PhraseCandidate) -> float:
weight
for enabled, weight in (
(candidate.source_yake, 2.0),
(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),
@@ -738,10 +655,6 @@ def candidate_source_names(candidate: PhraseCandidate) -> list[str]:
names.append("raw_ngram")
if candidate.source_yake:
names.append("yake")
if candidate.source_spacy_ner:
names.append("spacy_ner")
if candidate.source_spacy_noun_chunk:
names.append("spacy_noun_chunk")
if candidate.source_capitalized:
names.append("capitalized")
if candidate.source_metadata:
@@ -754,16 +667,14 @@ def extract_phrase_candidates_for_book(
chapters: Sequence[str],
config: EbookSearchConfig,
*,
nlp: SpacyLanguage | None = None,
metadata: Mapping[str, object] | None = None,
) -> list[PhraseCandidate]:
"""Extract, score, and limit phrase candidates for one book.
Args:
book_text (str): Full book text used for most extraction sources.
chapters (Sequence[str]): Chapter-like text blocks used for spaCy and frequency counts.
chapters (Sequence[str]): Chapter-like text blocks used for 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.
Returns:
@@ -792,16 +703,6 @@ def extract_phrase_candidates_for_book(
len(yake_candidates),
(perf_counter() - yake_started_at) * 1000,
)
spacy_candidates: dict[str, PhraseCandidate] = {}
if nlp is not None:
spacy_started_at = perf_counter()
for chapter in chapters:
spacy_candidates = merge_candidate_sources(spacy_candidates, extract_spacy_candidates(chapter, nlp, config))
logger.info(
"ebook_phrase_candidate_extract_spacy_complete candidates=%s duration_ms=%.1f",
len(spacy_candidates),
(perf_counter() - spacy_started_at) * 1000,
)
capitalized_started_at = perf_counter()
capitalized = extract_capitalized_phrases(book_text, config)
logger.info(
@@ -811,7 +712,7 @@ def extract_phrase_candidates_for_book(
)
metadata_candidates = extract_metadata_candidates(metadata, config)
candidates = merge_candidate_sources(raw, yake_candidates, spacy_candidates, capitalized, metadata_candidates)
candidates = merge_candidate_sources(raw, yake_candidates, capitalized, metadata_candidates)
enriched_started_at = perf_counter()
# 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.
@@ -831,12 +732,11 @@ def extract_phrase_candidates_for_book(
: config.protected_phrase_max_candidates_per_book
]
logger.info(
"ebook_phrase_candidate_extract_complete raw=%s yake=%s spacy=%s capitalized=%s metadata=%s "
"ebook_phrase_candidate_extract_complete raw=%s yake=%s capitalized=%s metadata=%s "
"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),
len(capitalized),
len(metadata_candidates),
pre_filter_count,
@@ -4,11 +4,11 @@ from __future__ import annotations
import asyncio
import logging
from collections import deque
from time import perf_counter
from typing import TYPE_CHECKING
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book
from python.ebook_search.protected_phrases.models import (
@@ -16,65 +16,93 @@ 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.pool import 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,
metadata_for_source_id,
new_candidate_row,
prune_unstorable_unjudged_candidate_phrases,
)
from python.orm.richie import EbookCandidatePhrase, EbookSource
from python.orm.common import get_async_postgres_engine
from python.orm.richie import EbookSource
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from concurrent.futures import Future
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.ext.asyncio import AsyncEngine
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.extraction import SpacyLanguage
from python.ebook_search.protected_phrases.models import PhraseCandidate
logger = logging.getLogger(__name__)
class BookHasNoChaptersError(ValueError):
"""Raised when a book has no indexed chapter text to generate phrases from."""
async def generate_candidate_phrases_for_books(
session: AsyncSession,
engine: AsyncEngine,
config: EbookSearchConfig,
*,
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.
Every book is submitted to the shared process pool up front and runs in parallel across the
pool's workers; the call blocks until all books have finished. Each worker opens its own
database engine from environment variables, loads the book's chapters, and commits the
book's candidates independently.
Args:
session (Session): Active database session.
engine (AsyncEngine): Engine used to read the book list in this process.
config (EbookSearchConfig): Runtime phrase-tuning settings.
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.
Results are collected in book order while the pool keeps working. A book failure (including
a book with no indexed chapters) is logged and counted as not built; the remaining books
are unaffected.
"""
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)
async with AsyncSession(engine, expire_on_commit=False) as session:
source_query = select(EbookSource.id).order_by(EbookSource.id)
source_ids = (await session.scalars(source_query)).all()
books_seen = len(source_ids)
logger.info(
"ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s max_candidates_per_book=%s",
"ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s "
"max_candidates_per_book=%s",
books_seen,
config.phrase_min_tokens,
config.phrase_max_tokens,
config.protected_phrase_max_candidates_per_book,
)
outcomes = await generate_candidates_for_sources_pooled(session, sources, config)
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
wrapped_futures = [
(
source_id,
asyncio.wrap_future(pool.submit(generate_candidate_phrases_for_book_in_worker, source_id, None, config)),
)
for source_id in source_ids
]
outcomes: list[BookCandidateResult] = []
for source_id, wrapped_future in wrapped_futures:
await asyncio.wait([wrapped_future])
exception = wrapped_future.exception()
if exception is not None:
logger.error(
"ebook_candidate_phrase_generation_book_failed source_id=%s",
source_id,
exc_info=exception,
)
outcomes.append(BookCandidateResult())
continue
saved_count = wrapped_future.result()
logger.info(
"ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s",
source_id,
saved_count,
)
outcomes.append(BookCandidateResult(candidates=saved_count, built=True))
result = PhraseCandidateGenerationResult(
books_seen=books_seen,
@@ -90,113 +118,27 @@ async def generate_candidate_phrases_for_books(
return result
async def generate_candidates_for_sources_pooled(
session: AsyncSession,
sources: Sequence[EbookSource],
config: EbookSearchConfig,
) -> list[BookCandidateResult]:
"""Generate candidate phrases for many books, extracting them concurrently in worker processes.
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.
sources (Sequence[EbookSource]): Indexed books to generate candidates for.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
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()
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,
saved_count,
(perf_counter() - book_started_at) * 1000,
)
return BookCandidateResult(candidates=saved_count, built=True)
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.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session; deletion and regeneration commit on it.
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.
Raises:
BookHasNoChaptersError: If the book has no indexed chapters. The deletion is rolled
back, so the book's existing phrases stay intact.
The deletion and regeneration share the caller's session, so they commit together; a
regeneration failure rolls the deletion back.
"""
started_at = perf_counter()
logger.info(
@@ -204,37 +146,14 @@ async def recalculate_candidate_phrases_for_book(
source.id,
source.title,
)
try:
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)
await session.commit()
return PhraseRecalculationResult(
book_id=source.id,
deleted_candidates=deleted.deleted_candidates,
deleted_protected_phrases=deleted.deleted_protected_phrases,
deleted_aliases=deleted.deleted_aliases,
deleted_mentions=deleted.deleted_mentions,
candidate_phrases=0,
)
candidate_count = await generate_candidate_phrases_for_book(
session,
source.id,
series_id=None,
chapters=chapters,
config=config,
nlp=nlp,
metadata=metadata_for_source(source),
replace_all=True,
use_process_pool=use_process_pool,
)
await session.commit()
except Exception:
await session.rollback()
logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id)
raise
result = PhraseRecalculationResult(
book_id=source.id,
@@ -258,57 +177,96 @@ async def recalculate_candidate_phrases_for_book(
return result
async def generate_candidate_phrases_for_book(
session: AsyncSession,
def generate_candidate_phrases_for_book_in_worker(
book_id: int,
series_id: int | None,
chapters: Sequence[str],
config: EbookSearchConfig,
*,
nlp: SpacyLanguage | None = None,
metadata: Mapping[str, object] | None = None,
replace_all: bool = False,
use_process_pool: bool = False,
) -> int:
"""Extract and store candidate phrases for one book without LLM judging.
"""Run one book's candidate generation in a pooled worker process.
The worker has no engine or session to inherit (neither can cross process boundaries), so
it creates its own engine from environment variables, opens the book's session on it, and
disposes the engine once the book is stored.
Args:
session (Session): Active database session.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope for the stored candidates.
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:
int: Number of candidate phrase rows stored.
"""
async def generate_with_worker_engine() -> int:
engine = get_async_postgres_engine(name="RICHIE", vector_engine=True, pool_size=1)
try:
async with AsyncSession(engine, expire_on_commit=False) as session:
return await generate_candidate_phrases_for_book(
session,
book_id,
series_id,
config,
)
finally:
await engine.dispose()
return asyncio.run(generate_with_worker_engine())
async def generate_candidate_phrases_for_book(
session: AsyncSession,
book_id: int,
series_id: int | None,
config: EbookSearchConfig,
*,
replace_all: bool = False,
) -> int:
"""Load a book's chapters and metadata, extract candidate phrases, and store them without LLM judging.
The session commits only when the whole book succeeds; any failure rolls the session back,
which also restores rows the caller deleted in the same transaction (e.g. a recalculation).
Args:
session (AsyncSession): Active database session; committed on success, rolled back on failure.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope for the stored candidates.
config (EbookSearchConfig): Runtime phrase-tuning settings.
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.
Returns:
int: Number of candidate phrase rows stored.
Raises:
BookHasNoChaptersError: If the book has no indexed chapter text.
"""
started_at = perf_counter()
chapters = await load_book_chapter_texts(session, book_id)
if not chapters:
await session.rollback()
message = f"book {book_id} has no indexed chapters"
raise BookHasNoChaptersError(message)
metadata = await metadata_for_source_id(session, book_id)
try:
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(
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,
candidates,
config,
replace_all=replace_all,
)
await session.commit()
except Exception:
await session.rollback()
raise
logger.info(
"ebook_candidate_phrase_generation_book_duration book_id=%s candidates=%s duration_ms=%.1f",
book_id,
@@ -19,11 +19,8 @@ class PhraseCandidate:
token_count (int): Number of normalized tokens in the phrase.
source_raw_ngram (bool): Whether the raw n-gram extractor produced the phrase.
source_yake (bool): Whether YAKE keyword extraction produced the phrase.
source_spacy_ner (bool): Whether spaCy named-entity recognition produced the phrase.
source_spacy_noun_chunk (bool): Whether spaCy noun chunking produced the phrase.
source_capitalized (bool): Whether the capitalized-run extractor produced the phrase.
source_metadata (bool): Whether book metadata produced the phrase.
spacy_label (str | None): spaCy entity label when NER produced the phrase.
raw_count (int): Occurrences counted across the book text.
chapter_count (int): Number of chapters containing the phrase.
yake_score (float | None): Raw YAKE score when available; lower is better.
@@ -36,11 +33,8 @@ class PhraseCandidate:
token_count: int
source_raw_ngram: bool = False
source_yake: bool = False
source_spacy_ner: bool = False
source_spacy_noun_chunk: bool = False
source_capitalized: bool = False
source_metadata: bool = False
spacy_label: str | None = None
raw_count: int = 0
chapter_count: int = 0
yake_score: float | None = None
@@ -9,21 +9,11 @@ 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__)
@@ -66,36 +56,3 @@ def shutdown_extraction_pool() -> 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)
@@ -295,11 +295,8 @@ def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
token_count=row.token_count,
source_raw_ngram=row.source_raw_ngram,
source_yake=row.source_yake,
source_spacy_ner=row.source_spacy_ner,
source_spacy_noun_chunk=row.source_spacy_noun_chunk,
source_capitalized=row.source_capitalized,
source_metadata=row.source_metadata,
spacy_label=row.spacy_label,
raw_count=row.raw_count,
chapter_count=row.chapter_count,
yake_score=row.yake_score,
@@ -334,11 +331,8 @@ def candidate_row_values(
"token_count": candidate.token_count,
"source_raw_ngram": candidate.source_raw_ngram,
"source_yake": candidate.source_yake,
"source_spacy_ner": candidate.source_spacy_ner,
"source_spacy_noun_chunk": candidate.source_spacy_noun_chunk,
"source_capitalized": candidate.source_capitalized,
"source_metadata": candidate.source_metadata,
"spacy_label": candidate.spacy_label,
"raw_count": candidate.raw_count,
"chapter_count": candidate.chapter_count,
"yake_score": candidate.yake_score,
@@ -460,11 +454,8 @@ def new_candidate_row(book_id: int, series_id: int | None, candidate: PhraseCand
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
-3
View File
@@ -167,11 +167,8 @@ class EbookCandidatePhrase(TableBase):
token_count: Mapped[int]
source_raw_ngram: Mapped[bool] = mapped_column(default=False)
source_yake: Mapped[bool] = mapped_column(default=False)
source_spacy_ner: Mapped[bool] = mapped_column(default=False)
source_spacy_noun_chunk: Mapped[bool] = mapped_column(default=False)
source_capitalized: Mapped[bool] = mapped_column(default=False)
source_metadata: Mapped[bool] = mapped_column(default=False)
spacy_label: Mapped[str | None]
raw_count: Mapped[int] = mapped_column(default=0)
chapter_count: Mapped[int] = mapped_column(default=0)
yake_score: Mapped[float | None]
+151 -19
View File
@@ -2,15 +2,16 @@
from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor
from datetime import UTC, datetime
from typing import TYPE_CHECKING
import pytest
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
from sqlalchemy.pool import StaticPool
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases import generate_ngrams
from python.ebook_search.protected_phrases.config import (
get_bad_ends,
get_most_common_words,
@@ -55,25 +56,42 @@ from python.orm.richie import (
)
if TYPE_CHECKING:
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, Generator
from pathlib import Path
from pytest_mock import MockerFixture
@pytest.fixture
async def engine() -> AsyncGenerator[AsyncEngine]:
"""Create a shared in-memory async database engine for phrase tests."""
test_engine = create_async_engine(
"sqlite+aiosqlite:///:memory:",
connect_args={"check_same_thread": False},
poolclass=StaticPool,
)
async def engine(tmp_path: Path) -> AsyncGenerator[AsyncEngine]:
"""Create a file-backed async database engine that worker threads can also reach."""
test_engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'phrases.db'}")
async with test_engine.begin() as connection:
await connection.run_sync(RichieBase.metadata.create_all)
yield test_engine
await test_engine.dispose()
@pytest.fixture
def worker_pool(engine: AsyncEngine, mocker: MockerFixture) -> Generator[ThreadPoolExecutor]:
"""Run pooled candidate generation in threads against the test database.
Spawned worker processes can see neither the test database nor test patches, so the shared
extraction pool is replaced with a thread pool and worker engines are built for the test
database instead of from Postgres environment variables.
"""
thread_pool = ThreadPoolExecutor(max_workers=1)
database_url = engine.url.render_as_string(hide_password=False)
mocker.patch.object(generate_ngrams, "get_extraction_pool", return_value=thread_pool)
mocker.patch.object(
generate_ngrams,
"get_async_postgres_engine",
side_effect=lambda **_kwargs: create_async_engine(database_url),
)
yield thread_pool
thread_pool.shutdown(wait=True)
@pytest.fixture
async def session(engine: AsyncEngine) -> AsyncGenerator[AsyncSession]:
"""Provide a session on the shared in-memory database."""
@@ -271,7 +289,9 @@ async def test_index_chunk_phrase_mentions_uses_normalized_window_lookup(
}
@pytest.mark.usefixtures("worker_pool")
async def test_generate_candidate_phrases_for_books_stores_unjudged_candidates(
engine: AsyncEngine,
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
@@ -300,8 +320,7 @@ async def test_generate_candidate_phrases_for_books_stores_unjudged_candidates(
}
)
result = await generate_candidate_phrases_for_books(session, build_config)
await session.commit()
result = await generate_candidate_phrases_for_books(engine, build_config)
candidate = await session.scalar(select(EbookCandidatePhrase))
assert result.books_seen == 1
@@ -313,7 +332,9 @@ async def test_generate_candidate_phrases_for_books_stores_unjudged_candidates(
assert await session.scalar(select(EbookProtectedPhrase)) is None
@pytest.mark.usefixtures("worker_pool")
async def test_generate_candidate_phrases_for_books_filters_one_token_and_one_use_candidates(
engine: AsyncEngine,
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
@@ -377,8 +398,7 @@ async def test_generate_candidate_phrases_for_books_filters_one_token_and_one_us
}
)
result = await generate_candidate_phrases_for_books(session, build_config)
await session.commit()
result = await generate_candidate_phrases_for_books(engine, build_config)
candidates = list(await session.scalars(select(EbookCandidatePhrase)))
phrase_norms = {candidate.phrase_norm for candidate in candidates}
@@ -394,7 +414,9 @@ async def test_generate_candidate_phrases_for_books_filters_one_token_and_one_us
assert all(not all(token in common_words for token in candidate.phrase_norm.split()) for candidate in candidates)
@pytest.mark.usefixtures("worker_pool")
async def test_generate_candidate_phrases_for_books_commits_after_each_book(
engine: AsyncEngine,
session: AsyncSession,
config: EbookSearchConfig,
mocker: MockerFixture,
@@ -431,7 +453,7 @@ async def test_generate_candidate_phrases_for_books_commits_after_each_book(
]
)
await session.commit()
commit_spy = mocker.spy(session, "commit")
commit_spy = mocker.spy(AsyncSession, "commit")
build_config = config.model_copy(
update={
"protected_phrase_max_candidates_per_book": 1,
@@ -440,13 +462,91 @@ async def test_generate_candidate_phrases_for_books_commits_after_each_book(
}
)
result = await generate_candidate_phrases_for_books(session, build_config)
result = await generate_candidate_phrases_for_books(engine, build_config)
assert result.books_built == 2
assert result.candidate_phrases == 2
assert commit_spy.call_count == 2
@pytest.mark.usefixtures("worker_pool")
async def test_generate_candidate_phrases_for_books_failure_keeps_committed_books(
engine: AsyncEngine,
session: AsyncSession,
config: EbookSearchConfig,
mocker: MockerFixture,
) -> None:
"""A failing book should be logged and skipped while the other books stay committed."""
first = await add_source(session)
second = await add_source(session, file_path="/library/book-2.epub", file_sha256="z" * 64)
session.add_all(
[
EbookChunk(
id=1,
source_id=first.id,
chapter_id=None,
chunk_index=0,
text="lock in lock in lock in",
token_start=0,
token_count=6,
page_label=None,
content_sha256="f" * 64,
search_text="lock in lock in lock in",
),
EbookChunk(
id=2,
source_id=second.id,
chapter_id=None,
chunk_index=0,
text="mage king mage king mage king",
token_start=0,
token_count=6,
page_label=None,
content_sha256="g" * 64,
search_text="mage king mage king mage king",
),
]
)
await session.commit()
build_config = config.model_copy(
update={
"protected_phrase_max_candidates_per_book": 1,
"phrase_min_tokens": 2,
"phrase_max_tokens": 2,
}
)
real_store = generate_ngrams.store_candidate_phrases_for_book
first_book_id = first.id
second_book_id = second.id
async def store_failing_second_book(
store_session: AsyncSession,
book_id: int,
series_id: int | None,
limited_candidates: list[PhraseCandidate],
store_config: EbookSearchConfig,
*,
replace_all: bool = False,
) -> int:
if book_id == second_book_id:
message = "storage exploded"
raise RuntimeError(message)
return await real_store(
store_session, book_id, series_id, limited_candidates, store_config, replace_all=replace_all
)
mocker.patch.object(generate_ngrams, "store_candidate_phrases_for_book", side_effect=store_failing_second_book)
result = await generate_candidate_phrases_for_books(engine, build_config)
stored_book_ids = set((await session.scalars(select(EbookCandidatePhrase.book_id))).all())
assert stored_book_ids == {first_book_id}
assert result.books_seen == 2
assert result.books_built == 1
assert result.candidate_phrases == 1
@pytest.mark.usefixtures("worker_pool")
async def test_judge_candidate_phrases_for_books_promotes_stored_candidates(
engine: AsyncEngine,
session: AsyncSession,
@@ -480,8 +580,7 @@ async def test_judge_candidate_phrases_for_books_promotes_stored_candidates(
"phrase_judge_phrase_workers": 1,
}
)
await generate_candidate_phrases_for_books(session, build_config)
await session.commit()
await generate_candidate_phrases_for_books(engine, build_config)
mocker.patch(
"python.ebook_search.protected_phrases.judge_ngrams.judge_candidate_async",
return_value=LLMJudgment(
@@ -513,6 +612,7 @@ async def test_judge_candidate_phrases_for_books_promotes_stored_candidates(
assert mention.phrase_id == phrase.id
@pytest.mark.usefixtures("worker_pool")
async def test_judge_candidate_phrases_for_books_logs_and_continues_after_book_failure(
engine: AsyncEngine,
session: AsyncSession,
@@ -561,8 +661,7 @@ async def test_judge_candidate_phrases_for_books_logs_and_continues_after_book_f
"phrase_judge_phrase_workers": 1,
}
)
await generate_candidate_phrases_for_books(session, build_config)
await session.commit()
await generate_candidate_phrases_for_books(engine, build_config)
def judge_or_fail(_client: object, _config: EbookSearchConfig, candidate: PhraseCandidate) -> LLMJudgment:
if candidate.phrase_norm == "lock in":
@@ -805,6 +904,7 @@ async def test_recalculate_candidate_phrases_for_book_removes_old_phrase_data(
result = await recalculate_candidate_phrases_for_book(session, source, build_config)
session.expire_all()
candidates = list(await session.scalars(select(EbookCandidatePhrase)))
assert result.deleted_candidates == 1
assert result.deleted_protected_phrases == 1
@@ -817,6 +917,38 @@ async def test_recalculate_candidate_phrases_for_book_removes_old_phrase_data(
assert await session.scalar(select(EbookChunkPhraseMention)) is None
async def test_recalculate_candidate_phrases_for_book_aborts_without_chapters(
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
"""Recalculating a book with no indexed chapters should raise and leave phrase data intact."""
source = await add_source(session)
existing_candidate = EbookCandidatePhrase(
book_id=source.id,
series_id=None,
phrase_text="old phrase",
phrase_norm="old phrase",
token_count=2,
source_raw_ngram=True,
raw_count=1,
chapter_count=1,
candidate_score=1.0,
llm_judged=False,
)
session.add(existing_candidate)
await session.flush()
existing_phrase = await add_phrase(session, source.id, phrase_text="old phrase", phrase_norm="old phrase")
await session.commit()
existing_candidate_id = existing_candidate.id
existing_phrase_id = existing_phrase.id
with pytest.raises(ValueError, match="no indexed chapters"):
await recalculate_candidate_phrases_for_book(session, source, config)
assert await session.scalar(select(EbookCandidatePhrase.id)) == existing_candidate_id
assert await session.scalar(select(EbookProtectedPhrase.id)) == existing_phrase_id
async def test_corpus_phrase_stats_counts_phrases_and_book_coverage(session: AsyncSession) -> None:
"""Corpus stats should count phrases plus how many books are generated and fully judged."""
unjudged_book = await add_source(session)
+1 -27
View File
@@ -567,33 +567,8 @@ def test_admin_page_shows_protected_phrase_stats(mocker: MockerFixture) -> None:
assert value in response.text
def test_ui_add_missing_phrases_generates_only_missing_books(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
def fake_generate(_session, _config, *, only_missing=False):
captured["only_missing"] = only_missing
return PhraseCandidateGenerationResult(books_seen=3, books_built=2, candidate_phrases=42)
mocker.patch(
"python.ebook_search.api.routes.admin.generate_candidate_phrases_for_books",
side_effect=fake_generate,
)
patch_app_runtime(mocker)
app = create_app()
with TestClient(app) as client:
response = client.post("/admin/phrases/generate-missing")
assert response.status_code == 200
assert captured["only_missing"] is True
assert "42 candidates stored" in response.text
def test_ui_regenerate_all_phrases_generates_every_book(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
def fake_generate(_session, _config, *, only_missing=False):
captured["only_missing"] = only_missing
def fake_generate(_session, _config):
return PhraseCandidateGenerationResult(books_seen=5, books_built=5, candidate_phrases=99)
mocker.patch(
@@ -607,7 +582,6 @@ def test_ui_regenerate_all_phrases_generates_every_book(mocker: MockerFixture) -
response = client.post("/admin/phrases/generate-all")
assert response.status_code == 200
assert captured["only_missing"] is False
assert "5 of 5 books" in response.text
-5
View File
@@ -1,5 +0,0 @@
{
home.sessionPath = [
"/home/richie/app_images/"
];
}
@@ -1,6 +1,5 @@
{
imports = [
../home/app_image_path.nix
../home/global.nix
../home/gui
];