Richie
1841ac69fa
treefmt / nix fmt (pull_request) Failing after 5s
pytest / pytest (pull_request) Failing after 36s
build_systems / build-brain (pull_request) Successful in 51s
build_systems / build-bob (pull_request) Successful in 55s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m30s
build_systems / build-jeeves (pull_request) Successful in 3m9s
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.
1090 lines
37 KiB
Python
1090 lines
37 KiB
Python
"""Tests for protected phrase extraction and matching."""
|
|
|
|
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 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,
|
|
)
|
|
from python.ebook_search.protected_phrases.extraction import (
|
|
extract_capitalized_phrases,
|
|
is_junk_phrase,
|
|
score_candidate,
|
|
)
|
|
from python.ebook_search.protected_phrases.generate_ngrams import (
|
|
generate_candidate_phrases_for_books,
|
|
recalculate_candidate_phrases_for_book,
|
|
)
|
|
from python.ebook_search.protected_phrases.judge_ngrams import (
|
|
judge_candidate_phrases_for_books,
|
|
prepare_book_judgment,
|
|
)
|
|
from python.ebook_search.protected_phrases.matching import (
|
|
detect_protected_phrases_for_query,
|
|
index_chunk_phrase_mentions,
|
|
load_phrase_lookup,
|
|
phrase_hit_counts_for_chunks,
|
|
phrase_hits_for_chunks,
|
|
resolve_overlaps,
|
|
)
|
|
from python.ebook_search.protected_phrases.models import (
|
|
ChunkPhraseHit,
|
|
HydratedPhraseMatch,
|
|
LLMJudgment,
|
|
PhraseCandidate,
|
|
)
|
|
from python.ebook_search.protected_phrases.store import book_ids_pending_first_judgment, corpus_phrase_stats
|
|
from python.ebook_search.protected_phrases.text_normalization import normalize_text, tokenize
|
|
from python.orm.richie import (
|
|
EbookCandidatePhrase,
|
|
EbookChunk,
|
|
EbookChunkPhraseMention,
|
|
EbookPhraseAlias,
|
|
EbookProtectedPhrase,
|
|
EbookSource,
|
|
RichieBase,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncGenerator, Generator
|
|
from pathlib import Path
|
|
|
|
from pytest_mock import MockerFixture
|
|
|
|
|
|
@pytest.fixture
|
|
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."""
|
|
async with AsyncSession(engine, expire_on_commit=False) as test_session:
|
|
yield test_session
|
|
|
|
|
|
@pytest.fixture
|
|
def config() -> EbookSearchConfig:
|
|
"""Provide default phrase-tuning settings for tests."""
|
|
return EbookSearchConfig()
|
|
|
|
|
|
def test_normalize_text_preserves_phrase_stopwords_and_word_order() -> None:
|
|
"""Normalization should not collapse protected phrases by removing stopwords."""
|
|
assert normalize_text("Haden\u2019s syndrome -- lock-in!") == "haden's syndrome lock in"
|
|
assert tokenize("House of the Dragon") == ["house", "of", "the", "dragon"]
|
|
|
|
|
|
def test_capitalized_phrase_keeps_bad_end_tokens_for_downranking(config: EbookSearchConfig) -> None:
|
|
"""Capitalized extraction should keep dangling endings so scoring can downrank them."""
|
|
candidates = extract_capitalized_phrases("Damion Montgomery and the left.", config)
|
|
|
|
assert "damion montgomery and the" in candidates
|
|
assert candidates["damion montgomery and the"].phrase_text == "Damion Montgomery and the"
|
|
assert candidates["damion montgomery and the"].phrase_norm.split()[-1] in get_bad_ends()
|
|
|
|
|
|
def test_is_junk_phrase_rejects_dialogue_verbs_contractions_and_majority_common() -> None:
|
|
"""Phrases the judge never keeps should be caught lexically before any LLM call."""
|
|
assert is_junk_phrase(["vann", "said"])
|
|
assert is_junk_phrase(["shook", "her", "head"])
|
|
assert is_junk_phrase(["i'm", "going"])
|
|
assert is_junk_phrase(["don't", "know"])
|
|
assert is_junk_phrase(["of", "the", "vault"])
|
|
|
|
|
|
def test_is_junk_phrase_keeps_possessives_and_half_common_world_terms() -> None:
|
|
"""Proper-noun possessives and phrasal world terms must survive the junk filter."""
|
|
assert not is_junk_phrase(["chapman's", "death"])
|
|
assert not is_junk_phrase(["boston", "bays'"])
|
|
assert not is_junk_phrase(["lock", "in"])
|
|
assert not is_junk_phrase(["data", "feed"])
|
|
assert not is_junk_phrase(["haden's", "syndrome"])
|
|
|
|
|
|
def test_score_candidate_rewards_multiple_non_raw_sources(config: EbookSearchConfig) -> None:
|
|
"""A second non-raw source should add its weight plus the multi-source bonus."""
|
|
single_source = PhraseCandidate(
|
|
phrase_text="lock in",
|
|
phrase_norm="lock in",
|
|
token_count=2,
|
|
source_capitalized=True,
|
|
raw_count=3,
|
|
chapter_count=2,
|
|
)
|
|
multi_source = PhraseCandidate(
|
|
phrase_text="lock in",
|
|
phrase_norm="lock in",
|
|
token_count=2,
|
|
source_capitalized=True,
|
|
source_yake=True,
|
|
raw_count=3,
|
|
chapter_count=2,
|
|
)
|
|
|
|
assert score_candidate(multi_source, config) == score_candidate(single_source, config) + 2.0 + 2.0
|
|
|
|
|
|
def test_score_candidate_caps_frequency_contribution(config: EbookSearchConfig) -> None:
|
|
"""Very frequent raw-only phrases should no longer out-score sourced entities."""
|
|
frequent_raw_only = PhraseCandidate(
|
|
phrase_text="mage king",
|
|
phrase_norm="mage king",
|
|
token_count=2,
|
|
source_raw_ngram=True,
|
|
raw_count=1000,
|
|
chapter_count=100,
|
|
)
|
|
|
|
assert score_candidate(frequent_raw_only, config) == 0.5 + 2.0 + 0.5
|
|
|
|
|
|
def test_score_candidate_weights_metadata_source(config: EbookSearchConfig) -> None:
|
|
"""Title and series phrases should get credit for the metadata source."""
|
|
metadata_only = PhraseCandidate(
|
|
phrase_text="lock in",
|
|
phrase_norm="lock in",
|
|
token_count=2,
|
|
source_metadata=True,
|
|
)
|
|
|
|
assert score_candidate(metadata_only, config) == 2.0 + 0.5
|
|
|
|
|
|
async def test_detect_protected_phrases_hydrates_alias_matches(
|
|
session: AsyncSession,
|
|
config: EbookSearchConfig,
|
|
) -> None:
|
|
"""Query detection should use RAM aliases and hydrate phrase metadata from the DB."""
|
|
source = await add_source(session)
|
|
phrase = await add_phrase(session, source.id, phrase_text="lock in", phrase_norm="lock in")
|
|
session.add(EbookPhraseAlias(phrase_id=phrase.id, alias_text="locked in", alias_norm="locked in"))
|
|
await session.commit()
|
|
|
|
matches = await detect_protected_phrases_for_query(session, "what is locked-in", config)
|
|
|
|
assert [(match.phrase_text, match.canonical_id, match.phrase_type) for match in matches] == [
|
|
("lock in", "condition:lock_in", "fictional_condition")
|
|
]
|
|
|
|
|
|
def test_resolve_overlaps_keeps_independent_nested_phrases() -> None:
|
|
"""Overlap resolution should keep useful nested concepts when metadata permits it."""
|
|
child = hydrated_match(
|
|
phrase_id=1,
|
|
phrase_text="mage king",
|
|
canonical_id="title:mage_king",
|
|
start_token=3,
|
|
end_token=5,
|
|
allow_nested=True,
|
|
)
|
|
parent = hydrated_match(
|
|
phrase_id=2,
|
|
phrase_text="mage king of mars",
|
|
canonical_id="entity:mage_king_of_mars",
|
|
start_token=3,
|
|
end_token=7,
|
|
suppress_children=False,
|
|
)
|
|
|
|
assert [match.phrase_text for match in resolve_overlaps([child, parent])] == [
|
|
"mage king of mars",
|
|
"mage king",
|
|
]
|
|
|
|
|
|
def test_resolve_overlaps_suppresses_weaker_same_canonical_match() -> None:
|
|
"""Same-canonical overlaps should keep the stronger evidence."""
|
|
weak = hydrated_match(
|
|
phrase_id=1,
|
|
phrase_text="lock",
|
|
canonical_id="condition:lock_in",
|
|
start_token=2,
|
|
end_token=3,
|
|
importance=0.2,
|
|
)
|
|
strong = hydrated_match(
|
|
phrase_id=2,
|
|
phrase_text="lock in",
|
|
canonical_id="condition:lock_in",
|
|
start_token=2,
|
|
end_token=4,
|
|
importance=0.9,
|
|
)
|
|
|
|
assert [match.phrase_text for match in resolve_overlaps([weak, strong])] == ["lock in"]
|
|
|
|
|
|
async def test_index_chunk_phrase_mentions_uses_normalized_window_lookup(
|
|
session: AsyncSession,
|
|
config: EbookSearchConfig,
|
|
) -> None:
|
|
"""Chunk indexing should store mention rows without scanning all phrases at query time."""
|
|
source = await add_source(session)
|
|
phrase = await add_phrase(session, source.id, phrase_text="lock in", phrase_norm="lock in")
|
|
chunk = EbookChunk(
|
|
id=1,
|
|
source_id=source.id,
|
|
chapter_id=None,
|
|
chunk_index=0,
|
|
text="Victims experienced lock-in during the crisis.",
|
|
token_start=0,
|
|
token_count=7,
|
|
page_label=None,
|
|
content_sha256="b" * 64,
|
|
search_text="Victims experienced lock-in during the crisis.",
|
|
)
|
|
session.add(chunk)
|
|
await session.commit()
|
|
lookup = await load_phrase_lookup(session, config, book_id=source.id)
|
|
|
|
count = await index_chunk_phrase_mentions(session, chunk, lookup=lookup)
|
|
await session.commit()
|
|
|
|
mention = await session.scalar(select(EbookChunkPhraseMention))
|
|
assert count == 1
|
|
assert mention is not None
|
|
assert mention.chunk_id == chunk.id
|
|
assert mention.phrase_id == phrase.id
|
|
assert chunk.text[mention.start_char : mention.end_char] == "lock-in"
|
|
assert await phrase_hit_counts_for_chunks(session, chunk_ids=[chunk.id], phrase_ids=[phrase.id]) == {chunk.id: 1}
|
|
assert await phrase_hits_for_chunks(session, chunk_ids=[chunk.id], phrase_ids=[phrase.id]) == {
|
|
chunk.id: (ChunkPhraseHit(phrase_id=phrase.id, phrase_text="lock in", mention_count=1),)
|
|
}
|
|
|
|
|
|
@pytest.mark.usefixtures("worker_pool")
|
|
async def test_generate_candidate_phrases_for_books_stores_unjudged_candidates(
|
|
engine: AsyncEngine,
|
|
session: AsyncSession,
|
|
config: EbookSearchConfig,
|
|
) -> None:
|
|
"""Candidate generation should populate the per-book phrase list without calling the LLM."""
|
|
source = await add_source(session)
|
|
session.add(
|
|
EbookChunk(
|
|
id=1,
|
|
source_id=source.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="d" * 64,
|
|
search_text="lock in lock in lock in",
|
|
)
|
|
)
|
|
await session.commit()
|
|
build_config = config.model_copy(
|
|
update={
|
|
"protected_phrase_max_candidates_per_book": 1,
|
|
"phrase_min_tokens": 2,
|
|
"phrase_max_tokens": 2,
|
|
}
|
|
)
|
|
|
|
result = await generate_candidate_phrases_for_books(engine, build_config)
|
|
|
|
candidate = await session.scalar(select(EbookCandidatePhrase))
|
|
assert result.books_seen == 1
|
|
assert result.books_built == 1
|
|
assert result.candidate_phrases == 1
|
|
assert candidate is not None
|
|
assert candidate.phrase_norm == "lock in"
|
|
assert candidate.llm_judged is False
|
|
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:
|
|
"""Candidate generation should not persist one-token or single-use phrases."""
|
|
source = await add_source(session)
|
|
session.add(
|
|
EbookChunk(
|
|
id=1,
|
|
source_id=source.id,
|
|
chapter_id=None,
|
|
chunk_index=0,
|
|
text=(
|
|
"Damion walked away. Damion woke up. rare phrase appeared once. "
|
|
"and then and then. lock in lock in lock in."
|
|
),
|
|
token_start=0,
|
|
token_count=16,
|
|
page_label=None,
|
|
content_sha256="i" * 64,
|
|
search_text=(
|
|
"Damion walked away. Damion woke up. rare phrase appeared once. "
|
|
"and then and then. lock in lock in lock in."
|
|
),
|
|
)
|
|
)
|
|
session.add_all(
|
|
[
|
|
EbookCandidatePhrase(
|
|
book_id=source.id,
|
|
series_id=None,
|
|
phrase_text="Damion",
|
|
phrase_norm="damion",
|
|
token_count=1,
|
|
source_capitalized=True,
|
|
raw_count=2,
|
|
chapter_count=1,
|
|
candidate_score=10.0,
|
|
llm_judged=False,
|
|
),
|
|
EbookCandidatePhrase(
|
|
book_id=source.id,
|
|
series_id=None,
|
|
phrase_text="rare phrase",
|
|
phrase_norm="rare phrase",
|
|
token_count=2,
|
|
source_raw_ngram=True,
|
|
raw_count=1,
|
|
chapter_count=1,
|
|
candidate_score=9.0,
|
|
llm_judged=False,
|
|
),
|
|
]
|
|
)
|
|
await session.commit()
|
|
build_config = config.model_copy(
|
|
update={
|
|
"protected_phrase_max_candidates_per_book": 50,
|
|
"phrase_min_tokens": 2,
|
|
"phrase_max_tokens": 2,
|
|
"phrase_raw_ngram_min_count": 2,
|
|
}
|
|
)
|
|
|
|
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}
|
|
common_words = get_most_common_words()
|
|
assert result.candidate_phrases == len(candidates)
|
|
assert "lock in" in phrase_norms
|
|
assert "damion" not in phrase_norms
|
|
assert "rare phrase" not in phrase_norms
|
|
assert "and then" not in phrase_norms
|
|
assert {"and", "then"}.issubset(common_words)
|
|
assert all(candidate.token_count >= 2 for candidate in candidates)
|
|
assert all(candidate.raw_count >= 2 for candidate in candidates)
|
|
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,
|
|
) -> None:
|
|
"""Candidate generation should persist each completed book independently."""
|
|
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()
|
|
commit_spy = mocker.spy(AsyncSession, "commit")
|
|
build_config = config.model_copy(
|
|
update={
|
|
"protected_phrase_max_candidates_per_book": 1,
|
|
"phrase_min_tokens": 2,
|
|
"phrase_max_tokens": 2,
|
|
}
|
|
)
|
|
|
|
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,
|
|
config: EbookSearchConfig,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
"""Judging should work from existing candidate rows and retain the candidate history."""
|
|
source = await add_source(session)
|
|
session.add(
|
|
EbookChunk(
|
|
id=1,
|
|
source_id=source.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="e" * 64,
|
|
search_text="lock in lock in lock in",
|
|
)
|
|
)
|
|
await session.commit()
|
|
build_config = config.model_copy(
|
|
update={
|
|
"protected_phrase_max_candidates_per_book": 1,
|
|
"protected_phrase_llm_candidates_per_book": 1,
|
|
"phrase_min_tokens": 2,
|
|
"phrase_max_tokens": 2,
|
|
"phrase_judge_book_workers": 1,
|
|
"phrase_judge_phrase_workers": 1,
|
|
}
|
|
)
|
|
await generate_candidate_phrases_for_books(engine, build_config)
|
|
mocker.patch(
|
|
"python.ebook_search.protected_phrases.judge_ngrams.judge_candidate_async",
|
|
return_value=LLMJudgment(
|
|
keep=True,
|
|
canonical="lock in",
|
|
category="fictional_condition",
|
|
aliases=(),
|
|
confidence=0.95,
|
|
importance=0.9,
|
|
),
|
|
)
|
|
|
|
result = await judge_candidate_phrases_for_books(engine, build_config)
|
|
|
|
session.expire_all()
|
|
candidate = await session.scalar(select(EbookCandidatePhrase))
|
|
phrase = await session.scalar(select(EbookProtectedPhrase))
|
|
mention = await session.scalar(select(EbookChunkPhraseMention))
|
|
assert result.books_seen == 1
|
|
assert result.books_judged == 1
|
|
assert result.candidates_judged == 1
|
|
assert result.protected_phrases == 1
|
|
assert candidate is not None
|
|
assert candidate.llm_judged is True
|
|
assert candidate.llm_keep is True
|
|
assert phrase is not None
|
|
assert phrase.source_candidate_id == candidate.id
|
|
assert mention is not None
|
|
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,
|
|
config: EbookSearchConfig,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
"""A failing book should be rolled back and logged while later books are still judged and 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,
|
|
"protected_phrase_llm_candidates_per_book": 1,
|
|
"phrase_min_tokens": 2,
|
|
"phrase_max_tokens": 2,
|
|
"phrase_judge_book_workers": 1,
|
|
"phrase_judge_phrase_workers": 1,
|
|
}
|
|
)
|
|
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":
|
|
message = "llm judge unavailable"
|
|
raise RuntimeError(message)
|
|
return LLMJudgment(
|
|
keep=True,
|
|
canonical="mage king",
|
|
category="title",
|
|
aliases=(),
|
|
confidence=0.95,
|
|
importance=0.9,
|
|
)
|
|
|
|
mocker.patch(
|
|
"python.ebook_search.protected_phrases.judge_ngrams.judge_candidate_async",
|
|
side_effect=judge_or_fail,
|
|
)
|
|
|
|
second_book_id = second.id
|
|
|
|
result = await judge_candidate_phrases_for_books(engine, build_config)
|
|
|
|
session.expire_all()
|
|
assert result.books_seen == 2
|
|
assert result.books_judged == 1
|
|
assert result.books_failed == 1
|
|
assert result.protected_phrases == 1
|
|
phrase = await session.scalar(select(EbookProtectedPhrase))
|
|
assert phrase is not None
|
|
assert phrase.book_id == second_book_id
|
|
|
|
|
|
async def test_judge_candidate_phrases_for_books_skips_unstorable_existing_candidates(
|
|
engine: AsyncEngine,
|
|
session: AsyncSession,
|
|
config: EbookSearchConfig,
|
|
mocker: MockerFixture,
|
|
) -> None:
|
|
"""Old one-token or single-use candidate rows should not be judged or promoted."""
|
|
source = await add_source(session)
|
|
session.add(
|
|
EbookChunk(
|
|
id=1,
|
|
source_id=source.id,
|
|
chapter_id=None,
|
|
chunk_index=0,
|
|
text="Damion Damion rare phrase",
|
|
token_start=0,
|
|
token_count=4,
|
|
page_label=None,
|
|
content_sha256="j" * 64,
|
|
search_text="Damion Damion rare phrase",
|
|
)
|
|
)
|
|
session.add_all(
|
|
[
|
|
EbookCandidatePhrase(
|
|
book_id=source.id,
|
|
series_id=None,
|
|
phrase_text="Damion",
|
|
phrase_norm="damion",
|
|
token_count=1,
|
|
source_capitalized=True,
|
|
raw_count=2,
|
|
chapter_count=1,
|
|
candidate_score=10.0,
|
|
llm_judged=False,
|
|
),
|
|
EbookCandidatePhrase(
|
|
book_id=source.id,
|
|
series_id=None,
|
|
phrase_text="rare phrase",
|
|
phrase_norm="rare phrase",
|
|
token_count=2,
|
|
source_raw_ngram=True,
|
|
raw_count=1,
|
|
chapter_count=1,
|
|
candidate_score=9.0,
|
|
llm_judged=False,
|
|
),
|
|
]
|
|
)
|
|
await session.commit()
|
|
judge_mock = mocker.patch("python.ebook_search.protected_phrases.judge_ngrams.judge_candidate_async")
|
|
build_config = config.model_copy(
|
|
update={
|
|
"phrase_judge_book_workers": 1,
|
|
"phrase_judge_phrase_workers": 1,
|
|
}
|
|
)
|
|
|
|
result = await judge_candidate_phrases_for_books(engine, build_config)
|
|
|
|
assert result.books_judged == 0
|
|
assert result.candidates_judged == 0
|
|
assert result.protected_phrases == 0
|
|
judge_mock.assert_not_called()
|
|
assert await session.scalar(select(EbookProtectedPhrase)) is None
|
|
|
|
|
|
async def test_prepare_book_judgment_skips_stored_junk_and_rescores_stale_rows(
|
|
engine: AsyncEngine,
|
|
session: AsyncSession,
|
|
config: EbookSearchConfig,
|
|
) -> None:
|
|
"""Judgment selection should drop junk rows and rank by fresh scores, not stored ones."""
|
|
source = await add_source(session)
|
|
session.add(
|
|
EbookChunk(
|
|
id=1,
|
|
source_id=source.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="k" * 64,
|
|
search_text="lock in lock in lock in",
|
|
)
|
|
)
|
|
session.add_all(
|
|
[
|
|
EbookCandidatePhrase(
|
|
book_id=source.id,
|
|
series_id=None,
|
|
phrase_text="vann said",
|
|
phrase_norm="vann said",
|
|
token_count=2,
|
|
source_raw_ngram=True,
|
|
raw_count=50,
|
|
chapter_count=10,
|
|
candidate_score=10.0,
|
|
llm_judged=False,
|
|
),
|
|
EbookCandidatePhrase(
|
|
book_id=source.id,
|
|
series_id=None,
|
|
phrase_text="boring pair",
|
|
phrase_norm="boring pair",
|
|
token_count=2,
|
|
source_raw_ngram=True,
|
|
raw_count=50,
|
|
chapter_count=10,
|
|
candidate_score=9.0,
|
|
llm_judged=False,
|
|
),
|
|
EbookCandidatePhrase(
|
|
book_id=source.id,
|
|
series_id=None,
|
|
phrase_text="lock in",
|
|
phrase_norm="lock in",
|
|
token_count=2,
|
|
source_capitalized=True,
|
|
source_yake=True,
|
|
raw_count=3,
|
|
chapter_count=2,
|
|
candidate_score=1.0,
|
|
llm_judged=False,
|
|
),
|
|
]
|
|
)
|
|
await session.commit()
|
|
build_config = config.model_copy(
|
|
update={
|
|
"protected_phrase_llm_candidates_per_book": 2,
|
|
"phrase_min_tokens": 2,
|
|
"phrase_max_tokens": 2,
|
|
}
|
|
)
|
|
|
|
prepared = await prepare_book_judgment(engine, source.id, build_config)
|
|
|
|
assert prepared is not None
|
|
work_items, target_remaining = prepared
|
|
judged_norms = [candidate.phrase_norm for _, candidate in work_items]
|
|
assert judged_norms == ["lock in", "boring pair"]
|
|
assert target_remaining == build_config.phrase_target_protected_per_book
|
|
fresh_scores = [candidate.candidate_score for _, candidate in work_items]
|
|
assert fresh_scores == sorted(fresh_scores, reverse=True)
|
|
assert fresh_scores[0] > 1.0
|
|
|
|
|
|
async def test_recalculate_candidate_phrases_for_book_removes_old_phrase_data(
|
|
session: AsyncSession,
|
|
config: EbookSearchConfig,
|
|
) -> None:
|
|
"""Book-level recalculation should clear stale candidates, protected phrases, aliases, and mentions."""
|
|
source = await add_source(session)
|
|
session.add(
|
|
EbookChunk(
|
|
id=1,
|
|
source_id=source.id,
|
|
chapter_id=None,
|
|
chunk_index=0,
|
|
text="new phrase new phrase new phrase",
|
|
token_start=0,
|
|
token_count=6,
|
|
page_label=None,
|
|
content_sha256="h" * 64,
|
|
search_text="new phrase new phrase new phrase",
|
|
)
|
|
)
|
|
old_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(old_candidate)
|
|
await session.flush()
|
|
old_phrase = await add_phrase(session, source.id, phrase_text="old phrase", phrase_norm="old phrase")
|
|
old_phrase.source_candidate_id = old_candidate.id
|
|
session.add(EbookPhraseAlias(phrase_id=old_phrase.id, alias_text="old alias", alias_norm="old alias"))
|
|
session.add(
|
|
EbookChunkPhraseMention(
|
|
chunk_id=1,
|
|
phrase_id=old_phrase.id,
|
|
book_id=source.id,
|
|
series_id=None,
|
|
start_char=0,
|
|
end_char=10,
|
|
)
|
|
)
|
|
await session.commit()
|
|
build_config = config.model_copy(
|
|
update={
|
|
"protected_phrase_max_candidates_per_book": 1,
|
|
"phrase_min_tokens": 2,
|
|
"phrase_max_tokens": 2,
|
|
}
|
|
)
|
|
|
|
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
|
|
assert result.deleted_aliases == 1
|
|
assert result.deleted_mentions == 1
|
|
assert result.candidate_phrases == 1
|
|
assert [candidate.phrase_norm for candidate in candidates] == ["new phrase"]
|
|
assert await session.scalar(select(EbookProtectedPhrase)) is None
|
|
assert await session.scalar(select(EbookPhraseAlias)) is None
|
|
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)
|
|
judged_book = await add_source(session, file_path="/library/judged.epub", file_sha256="b" * 64)
|
|
mixed_book = await add_source(session, file_path="/library/mixed.epub", file_sha256="c" * 64)
|
|
await add_source(session, file_path="/library/empty.epub", file_sha256="d" * 64)
|
|
await add_candidate(session, unjudged_book.id, phrase_norm="rare phrase", llm_judged=False)
|
|
await add_candidate(session, judged_book.id, phrase_norm="lock in", llm_judged=True)
|
|
await add_candidate(session, mixed_book.id, phrase_norm="haden's syndrome", llm_judged=True)
|
|
await add_candidate(session, mixed_book.id, phrase_norm="boston bays", llm_judged=False)
|
|
await add_phrase(session, judged_book.id, phrase_text="Lock In", phrase_norm="lock in")
|
|
|
|
stats = await corpus_phrase_stats(session)
|
|
|
|
assert stats.total_books == 4
|
|
assert stats.books_with_candidates == 3
|
|
assert stats.books_fully_judged == 1
|
|
assert stats.candidate_phrases == 4
|
|
assert stats.judged_candidates == 2
|
|
assert stats.unjudged_candidates == 2
|
|
assert stats.protected_phrases == 1
|
|
|
|
|
|
async def test_book_ids_pending_first_judgment_returns_only_never_judged_books(session: AsyncSession) -> None:
|
|
"""Only books whose candidates are all unjudged should be pending a first judgment."""
|
|
unjudged_book = await add_source(session)
|
|
judged_book = await add_source(session, file_path="/library/judged.epub", file_sha256="b" * 64)
|
|
mixed_book = await add_source(session, file_path="/library/mixed.epub", file_sha256="c" * 64)
|
|
await add_candidate(session, unjudged_book.id, phrase_norm="rare phrase", llm_judged=False)
|
|
await add_candidate(session, judged_book.id, phrase_norm="lock in", llm_judged=True)
|
|
await add_candidate(session, mixed_book.id, phrase_norm="haden's syndrome", llm_judged=True)
|
|
await add_candidate(session, mixed_book.id, phrase_norm="boston bays", llm_judged=False)
|
|
|
|
assert await book_ids_pending_first_judgment(session) == [unjudged_book.id]
|
|
|
|
|
|
async def add_candidate(
|
|
session: AsyncSession,
|
|
book_id: int,
|
|
*,
|
|
phrase_norm: str,
|
|
llm_judged: bool,
|
|
) -> EbookCandidatePhrase:
|
|
"""Add a minimal candidate phrase row."""
|
|
candidate = EbookCandidatePhrase(
|
|
book_id=book_id,
|
|
series_id=None,
|
|
phrase_text=phrase_norm,
|
|
phrase_norm=phrase_norm,
|
|
token_count=len(phrase_norm.split()),
|
|
source_raw_ngram=True,
|
|
raw_count=3,
|
|
chapter_count=2,
|
|
candidate_score=5.0,
|
|
llm_judged=llm_judged,
|
|
)
|
|
session.add(candidate)
|
|
await session.flush()
|
|
return candidate
|
|
|
|
|
|
async def add_source(
|
|
session: AsyncSession,
|
|
*,
|
|
file_path: str = "/library/book.epub",
|
|
file_sha256: str = "a" * 64,
|
|
) -> EbookSource:
|
|
"""Add a minimal ebook source."""
|
|
source = EbookSource(
|
|
title="Book",
|
|
author="Author",
|
|
language=None,
|
|
publisher=None,
|
|
identifier=None,
|
|
file_path=file_path,
|
|
file_sha256=file_sha256,
|
|
file_mtime=datetime.now(tz=UTC),
|
|
file_size=10,
|
|
)
|
|
session.add(source)
|
|
await session.flush()
|
|
return source
|
|
|
|
|
|
async def add_phrase(
|
|
session: AsyncSession,
|
|
book_id: int,
|
|
*,
|
|
phrase_text: str,
|
|
phrase_norm: str,
|
|
) -> EbookProtectedPhrase:
|
|
"""Add a protected phrase row."""
|
|
phrase = EbookProtectedPhrase(
|
|
book_id=book_id,
|
|
series_id=None,
|
|
phrase_text=phrase_text,
|
|
phrase_norm=phrase_norm,
|
|
canonical_id="condition:lock_in",
|
|
phrase_type="fictional_condition",
|
|
token_count=len(phrase_norm.split()),
|
|
confidence=0.91,
|
|
importance=0.85,
|
|
allow_nested=False,
|
|
suppress_children=True,
|
|
source_candidate_id=None,
|
|
)
|
|
session.add(phrase)
|
|
await session.flush()
|
|
return phrase
|
|
|
|
|
|
def hydrated_match(
|
|
*,
|
|
phrase_id: int,
|
|
phrase_text: str,
|
|
canonical_id: str,
|
|
start_token: int,
|
|
end_token: int,
|
|
importance: float = 0.8,
|
|
allow_nested: bool = False,
|
|
suppress_children: bool = True,
|
|
) -> HydratedPhraseMatch:
|
|
"""Build a hydrated match for overlap tests."""
|
|
return HydratedPhraseMatch(
|
|
phrase_id=phrase_id,
|
|
matched_norm=phrase_text,
|
|
phrase_text=phrase_text,
|
|
phrase_norm=phrase_text,
|
|
canonical_id=canonical_id,
|
|
phrase_type=None,
|
|
token_count=end_token - start_token,
|
|
confidence=0.9,
|
|
importance=importance,
|
|
allow_nested=allow_nested,
|
|
suppress_children=suppress_children,
|
|
start_token=start_token,
|
|
end_token=end_token,
|
|
)
|