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

fix(protected-phrases): isolate phrase generation per book

Browse Source 
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.
This commit is contained in:
Richie
2026-07-24 11:38:50 -04:00
parent 29a51eb1b8
commit 58be234d7f
6 changed files with 300 additions and 306 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/ebook_search/test_protected_phrases.py
+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)
tests/ebook_search/test_ui.py
+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