Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
659b1e2798 | ||
|
|
b13bb39401 | ||
|
|
ffdb93d352 | ||
|
|
0cf05c52b9 |
@@ -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 ###
|
||||
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
];
|
||||
|
||||
Reference in New Issue
Block a user