test(ebook): cover protected phrases and migrate suite to async
Add test_protected_phrases.py covering phrase-matching behavior in the RAG engine, and update the existing ebook_search tests to use the async SQLAlchemy engine/session (create_async_engine, AsyncSession) and async HTTP paths.
This commit is contained in:
@@ -10,8 +10,8 @@ from types import ModuleType
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import create_engine, select
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
|
||||
|
||||
from python.ebook_search.answer import answer_query
|
||||
from python.ebook_search.bm25_corpus import (
|
||||
@@ -77,10 +77,17 @@ def test_reciprocal_rank_fusion_combines_vector_and_bm25_rankings() -> None:
|
||||
assert fused[0].fused_score == fused[0].score
|
||||
|
||||
|
||||
def test_find_existing_source_matches_path_or_hash() -> None:
|
||||
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||
RichieBase.metadata.create_all(engine)
|
||||
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
|
||||
async def build_async_engine() -> AsyncEngine:
|
||||
"""Create an in-memory async engine with the Richie schema."""
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
async with engine.begin() as connection:
|
||||
await connection.run_sync(RichieBase.metadata.create_all)
|
||||
return engine
|
||||
|
||||
|
||||
async def test_find_existing_source_matches_path_or_hash() -> None:
|
||||
engine = await build_async_engine()
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
source = EbookSource(
|
||||
title="Book",
|
||||
author=None,
|
||||
@@ -93,16 +100,15 @@ def test_find_existing_source_matches_path_or_hash() -> None:
|
||||
file_size=10,
|
||||
)
|
||||
session.add(source)
|
||||
session.commit()
|
||||
await session.commit()
|
||||
|
||||
assert find_existing_source(session, Path("/old/book.epub"), "b" * 64) == source
|
||||
assert find_existing_source(session, Path("/new/book.epub"), "a" * 64) == source
|
||||
assert await find_existing_source(session, Path("/old/book.epub"), "b" * 64) == source
|
||||
assert await find_existing_source(session, Path("/new/book.epub"), "a" * 64) == source
|
||||
|
||||
|
||||
def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> None:
|
||||
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||
RichieBase.metadata.create_all(engine)
|
||||
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
|
||||
async def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> None:
|
||||
engine = await build_async_engine()
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
source = EbookSource(
|
||||
title="Book",
|
||||
author="Author",
|
||||
@@ -115,10 +121,10 @@ def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> N
|
||||
file_size=10,
|
||||
)
|
||||
session.add(source)
|
||||
session.flush()
|
||||
await session.flush()
|
||||
chapter = EbookChapter(source_id=source.id, spine_index=0, title="Chapter", href=None)
|
||||
session.add(chapter)
|
||||
session.flush()
|
||||
await session.flush()
|
||||
session.add(
|
||||
EbookChunk(
|
||||
id=1,
|
||||
@@ -133,9 +139,9 @@ def test_bm25_corpus_uses_existing_search_text_without_duplicate_metadata() -> N
|
||||
search_text="Book Author Chapter content",
|
||||
)
|
||||
)
|
||||
session.commit()
|
||||
await session.commit()
|
||||
|
||||
records, texts = fetch_bm25_corpus_records(session)
|
||||
records, texts = await fetch_bm25_corpus_records(session)
|
||||
|
||||
assert texts == ["Book Author Chapter content"]
|
||||
assert records[0]["chunk_id"] == 1
|
||||
@@ -370,7 +376,7 @@ def test_load_bm25_corpus_raises_when_index_is_missing(mocker: MockerFixture, tm
|
||||
load_bm25_corpus.cache_clear()
|
||||
|
||||
|
||||
def test_ensure_bm25_corpus_refreshes_missing_index(mocker: MockerFixture) -> None:
|
||||
async def test_ensure_bm25_corpus_refreshes_missing_index(mocker: MockerFixture) -> None:
|
||||
refreshed: list[object] = []
|
||||
db_updated_at = datetime.now(tz=UTC)
|
||||
|
||||
@@ -385,12 +391,12 @@ def test_ensure_bm25_corpus_refreshes_missing_index(mocker: MockerFixture) -> No
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
session = object()
|
||||
|
||||
ensure_bm25_corpus(session, config)
|
||||
await ensure_bm25_corpus(session, config)
|
||||
|
||||
assert refreshed == [(session, config, db_updated_at)]
|
||||
|
||||
|
||||
def test_ensure_bm25_corpus_refreshes_stale_index(mocker: MockerFixture) -> None:
|
||||
async def test_ensure_bm25_corpus_refreshes_stale_index(mocker: MockerFixture) -> None:
|
||||
refreshed: list[object] = []
|
||||
created_at = datetime(2026, 1, 1, tzinfo=UTC)
|
||||
db_updated_at = datetime(2026, 1, 2, tzinfo=UTC)
|
||||
@@ -407,7 +413,7 @@ def test_ensure_bm25_corpus_refreshes_stale_index(mocker: MockerFixture) -> None
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
session = object()
|
||||
|
||||
ensure_bm25_corpus(session, config)
|
||||
await ensure_bm25_corpus(session, config)
|
||||
|
||||
assert refreshed == [(session, config, db_updated_at)]
|
||||
|
||||
@@ -420,14 +426,13 @@ def test_supported_embedding_models_match_service_names() -> None:
|
||||
}
|
||||
|
||||
|
||||
def test_ensure_embedding_models_registers_service_names() -> None:
|
||||
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||
RichieBase.metadata.create_all(engine)
|
||||
with sessionmaker(bind=engine, expire_on_commit=False, future=True)() as session:
|
||||
ensure_embedding_models(session)
|
||||
session.commit()
|
||||
async def test_ensure_embedding_models_registers_service_names() -> None:
|
||||
engine = await build_async_engine()
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
await ensure_embedding_models(session)
|
||||
await session.commit()
|
||||
|
||||
models = list(session.scalars(select(EbookEmbeddingModel).order_by(EbookEmbeddingModel.name)))
|
||||
models = list(await session.scalars(select(EbookEmbeddingModel).order_by(EbookEmbeddingModel.name)))
|
||||
|
||||
assert [(model.name, model.dimension) for model in models] == [
|
||||
("qwen3-embedding-0.6b", 1024),
|
||||
@@ -496,10 +501,10 @@ def test_chat_api_key_falls_back_to_ollama_api_key(mocker: MockerFixture) -> Non
|
||||
assert config.vllm_api_key == "ollama-key"
|
||||
|
||||
|
||||
def test_answer_query_does_not_call_model_when_disabled() -> None:
|
||||
async def test_answer_query_does_not_call_model_when_disabled(mocker: MockerFixture) -> None:
|
||||
config = load_config().model_copy(update={"answer_enabled": False})
|
||||
result = SearchResult(chunk_id=1, text="source text", source_title="Book")
|
||||
|
||||
answer = answer_query("question", [result], config)
|
||||
answer = await answer_query(mocker.Mock(), "question", [result], config)
|
||||
|
||||
assert "Answer generation is disabled" in answer
|
||||
|
||||
@@ -5,7 +5,7 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from python.ebook_search.api.main import create_app
|
||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||
@@ -66,8 +66,8 @@ def test_is_confident_against_threshold() -> None:
|
||||
|
||||
def patch_app_runtime(mocker: MockerFixture):
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.main.get_postgres_engine",
|
||||
side_effect=lambda **_kwargs: create_engine("sqlite+pysqlite:///:memory:", future=True),
|
||||
"python.ebook_search.api.main.get_async_postgres_engine",
|
||||
side_effect=lambda **_kwargs: create_async_engine("sqlite+aiosqlite:///:memory:"),
|
||||
)
|
||||
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
|
||||
|
||||
@@ -75,11 +75,12 @@ def patch_app_runtime(mocker: MockerFixture):
|
||||
def test_low_confidence_skips_answer_generation(mocker: MockerFixture) -> None:
|
||||
called = False
|
||||
|
||||
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
del phrase_matching
|
||||
return SearchResponse(query=query, rank_label="Hybrid", results=make_results(1, vector_score=0.05))
|
||||
|
||||
def fake_answer_query(_query, _results, _config):
|
||||
def fake_answer_query(_client, _query, _results, _config):
|
||||
nonlocal called
|
||||
called = True
|
||||
return "answer"
|
||||
@@ -104,14 +105,15 @@ def test_low_confidence_skips_answer_generation(mocker: MockerFixture) -> None:
|
||||
|
||||
|
||||
def test_invalid_citation_is_flagged(mocker: MockerFixture) -> None:
|
||||
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
del phrase_matching
|
||||
return SearchResponse(query=query, rank_label="Hybrid", results=make_results(2, vector_score=0.9))
|
||||
|
||||
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.search.answer_query",
|
||||
side_effect=lambda _query, _results, _config: "Per the text [9].",
|
||||
side_effect=lambda _client, _query, _results, _config: "Per the text [9].",
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
@@ -126,14 +128,15 @@ def test_invalid_citation_is_flagged(mocker: MockerFixture) -> None:
|
||||
|
||||
|
||||
def test_grounded_answer_has_no_warning_badge(mocker: MockerFixture) -> None:
|
||||
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
del phrase_matching
|
||||
return SearchResponse(query=query, rank_label="Hybrid", results=make_results(2, vector_score=0.9))
|
||||
|
||||
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.search.answer_query",
|
||||
side_effect=lambda _query, _results, _config: "Grounded in [1] and [2].",
|
||||
side_effect=lambda _client, _query, _results, _config: "Grounded in [1] and [2].",
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
|
||||
@@ -5,7 +5,7 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from python.ebook_search.api.main import create_app
|
||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||
@@ -18,18 +18,18 @@ if TYPE_CHECKING:
|
||||
|
||||
def fake_get_postgres_engine(**_kwargs):
|
||||
"""Return an in-memory engine for route tests."""
|
||||
return create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||
return create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
|
||||
|
||||
def patch_app_runtime(mocker: MockerFixture):
|
||||
mocker.patch("python.ebook_search.api.main.get_postgres_engine", side_effect=fake_get_postgres_engine)
|
||||
mocker.patch("python.ebook_search.api.main.get_async_postgres_engine", side_effect=fake_get_postgres_engine)
|
||||
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
|
||||
|
||||
|
||||
def patch_dependencies(mocker: MockerFixture, *, database=True, embedding=True, chat=True, bm25="ok"):
|
||||
mocker.patch(f"{HEALTH_MODULE}.check_database", side_effect=lambda _session: database)
|
||||
mocker.patch(f"{HEALTH_MODULE}.check_embedding_endpoint", side_effect=lambda _config: embedding)
|
||||
mocker.patch(f"{HEALTH_MODULE}.check_chat_endpoint", side_effect=lambda _config: chat)
|
||||
mocker.patch(f"{HEALTH_MODULE}.check_embedding_endpoint", side_effect=lambda _client, _config: embedding)
|
||||
mocker.patch(f"{HEALTH_MODULE}.check_chat_endpoint", side_effect=lambda _client, _config: chat)
|
||||
mocker.patch(f"{HEALTH_MODULE}.check_bm25_status", side_effect=lambda _config: bm25)
|
||||
|
||||
|
||||
|
||||
@@ -16,7 +16,14 @@ if TYPE_CHECKING:
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
|
||||
def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None:
|
||||
def make_async_client(mocker: MockerFixture, fake_post) -> httpx.AsyncClient:
|
||||
"""Build a mock async client whose post call is served by fake_post."""
|
||||
client = mocker.MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = mocker.AsyncMock(side_effect=fake_post)
|
||||
return client
|
||||
|
||||
|
||||
async def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_post(url: str, **kwargs: object) -> httpx.Response:
|
||||
@@ -28,7 +35,7 @@ def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None
|
||||
request=httpx.Request("POST", url),
|
||||
)
|
||||
|
||||
mocker.patch.object(httpx, "post", side_effect=fake_post)
|
||||
client = make_async_client(mocker, fake_post)
|
||||
config = EbookSearchConfig(
|
||||
rerank=RerankConfig(enabled=False),
|
||||
vllm_base_url="https://ollama.com/v1",
|
||||
@@ -36,7 +43,8 @@ def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None
|
||||
chat_model="deepseek-v4-flash",
|
||||
)
|
||||
|
||||
answer = answer_query("question", [SearchResult(chunk_id=1, text="source", source_title="Book")], config)
|
||||
results = [SearchResult(chunk_id=1, text="source", source_title="Book")]
|
||||
answer = await answer_query(client, "question", results, config)
|
||||
|
||||
assert answer == "grounded answer"
|
||||
assert captured["url"] == "https://ollama.com/v1/chat/completions"
|
||||
@@ -48,7 +56,7 @@ def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None
|
||||
assert payload["model"] == "deepseek-v4-flash"
|
||||
|
||||
|
||||
def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None:
|
||||
async def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
vector = [0.0] * 1024
|
||||
|
||||
@@ -61,14 +69,14 @@ def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None:
|
||||
request=httpx.Request("POST", url),
|
||||
)
|
||||
|
||||
mocker.patch.object(httpx, "post", side_effect=fake_post)
|
||||
client = make_async_client(mocker, fake_post)
|
||||
config = EbookSearchConfig(
|
||||
rerank=RerankConfig(enabled=False),
|
||||
embedding_base_url="http://bob:8000/v1",
|
||||
embedding_model="qwen3-embedding-0.6b",
|
||||
)
|
||||
|
||||
embeddings = embed_texts(["hello"], config)
|
||||
embeddings = await embed_texts(client, ["hello"], config)
|
||||
|
||||
assert embeddings == [vector]
|
||||
assert captured["url"] == "http://bob:8000/v1/embeddings"
|
||||
@@ -78,12 +86,12 @@ def test_embed_texts_uses_httpx_embeddings(mocker: MockerFixture) -> None:
|
||||
assert kwargs["json"] == {"model": "qwen3-embedding-0.6b", "input": ["hello"]}
|
||||
|
||||
|
||||
def test_embed_texts_rejects_bad_response_shape(mocker: MockerFixture) -> None:
|
||||
async def test_embed_texts_rejects_bad_response_shape(mocker: MockerFixture) -> None:
|
||||
def fake_post(url: str, **_kwargs: object) -> httpx.Response:
|
||||
return httpx.Response(200, json={"data": [{}]}, request=httpx.Request("POST", url))
|
||||
|
||||
mocker.patch.object(httpx, "post", side_effect=fake_post)
|
||||
client = make_async_client(mocker, fake_post)
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
|
||||
with pytest.raises(RuntimeError, match="Embedding request failed"):
|
||||
embed_texts(["hello"], config)
|
||||
await embed_texts(client, ["hello"], config)
|
||||
|
||||
@@ -0,0 +1,957 @@
|
||||
"""Tests for protected phrase extraction and matching."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
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.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
|
||||
|
||||
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 with test_engine.begin() as connection:
|
||||
await connection.run_sync(RichieBase.metadata.create_all)
|
||||
yield test_engine
|
||||
await test_engine.dispose()
|
||||
|
||||
|
||||
@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),)
|
||||
}
|
||||
|
||||
|
||||
async def test_generate_candidate_phrases_for_books_stores_unjudged_candidates(
|
||||
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(session, build_config)
|
||||
await session.commit()
|
||||
|
||||
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
|
||||
|
||||
|
||||
async def test_generate_candidate_phrases_for_books_filters_one_token_and_one_use_candidates(
|
||||
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(session, build_config)
|
||||
await session.commit()
|
||||
|
||||
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)
|
||||
|
||||
|
||||
async def test_generate_candidate_phrases_for_books_commits_after_each_book(
|
||||
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(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(session, build_config)
|
||||
|
||||
assert result.books_built == 2
|
||||
assert result.candidate_phrases == 2
|
||||
assert commit_spy.call_count == 2
|
||||
|
||||
|
||||
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(session, build_config)
|
||||
await session.commit()
|
||||
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
|
||||
|
||||
|
||||
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(session, build_config)
|
||||
await session.commit()
|
||||
|
||||
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)
|
||||
|
||||
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_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,
|
||||
)
|
||||
@@ -2,10 +2,11 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from threading import Event
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||
from python.ebook_search.search import SearchResult, search_ebooks
|
||||
@@ -14,18 +15,18 @@ if TYPE_CHECKING:
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
|
||||
def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -> None:
|
||||
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||
async def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -> None:
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
vector_started = Event()
|
||||
bm25_started = Event()
|
||||
received_engines: list[object] = []
|
||||
|
||||
def fake_vector_candidates(received_engine, query, _config):
|
||||
async def fake_vector_candidates(received_engine, _client, query, _config):
|
||||
"""Return vector candidates after confirming BM25 has started."""
|
||||
received_engines.append(received_engine)
|
||||
assert query == "what is parallel"
|
||||
vector_started.set()
|
||||
assert bm25_started.wait(timeout=2)
|
||||
assert await asyncio.to_thread(bm25_started.wait, 2)
|
||||
return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]
|
||||
|
||||
def fake_bm25_candidates(query, _config):
|
||||
@@ -39,7 +40,7 @@ def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -
|
||||
mocker.patch("python.ebook_search.search.bm25_candidates", side_effect=fake_bm25_candidates)
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
|
||||
response = search_ebooks(engine, "what is parallel", config)
|
||||
response = await search_ebooks(engine, mocker.Mock(), "what is parallel", config)
|
||||
|
||||
timings = {step.name: step for step in response.timings}
|
||||
assert [result.chunk_id for result in response.results] == [1, 2]
|
||||
@@ -47,3 +48,26 @@ def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -
|
||||
assert timings["BM25 search"].counts_toward_total is False
|
||||
assert timings["Hybrid retrieval"].counts_toward_total is True
|
||||
assert received_engines == [engine]
|
||||
|
||||
|
||||
async def test_search_ebooks_skips_phrase_matching_when_disabled(mocker: MockerFixture) -> None:
|
||||
"""Phrase matching can be disabled for one search request."""
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
mocker.patch(
|
||||
"python.ebook_search.search.vector_candidates",
|
||||
return_value=[SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)],
|
||||
)
|
||||
mocker.patch("python.ebook_search.search.bm25_candidates", return_value=[])
|
||||
detect_mock = mocker.patch("python.ebook_search.search.query_phrase_matches")
|
||||
boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
|
||||
response = await search_ebooks(engine, mocker.Mock(), "what is parallel", config, phrase_matching=False)
|
||||
|
||||
timing_names = {step.name for step in response.timings}
|
||||
assert [result.chunk_id for result in response.results] == [1]
|
||||
assert response.phrase_matches == ()
|
||||
assert "Protected phrase detection skipped" in timing_names
|
||||
assert "Phrase mention boost skipped" in timing_names
|
||||
detect_mock.assert_not_called()
|
||||
boost_mock.assert_not_called()
|
||||
|
||||
@@ -24,6 +24,13 @@ def candidates() -> list[SearchResult]:
|
||||
]
|
||||
|
||||
|
||||
def make_async_client(mocker: MockerFixture, fake_post) -> httpx.AsyncClient:
|
||||
"""Build a mock async client whose post call is served by fake_post."""
|
||||
client = mocker.MagicMock(spec=httpx.AsyncClient)
|
||||
client.post = mocker.AsyncMock(side_effect=fake_post)
|
||||
return client
|
||||
|
||||
|
||||
def rerank_response(payload: dict[str, object] | None = None, *, content: bytes | None = None) -> httpx.Response:
|
||||
return httpx.Response(
|
||||
200,
|
||||
@@ -59,7 +66,7 @@ def test_reranking_disabled_returns_original_fused_order() -> None:
|
||||
assert [result.chunk_id for result in response.results] == [1, 2]
|
||||
|
||||
|
||||
def test_reranking_enabled_reorders_candidates(mocker: MockerFixture) -> None:
|
||||
async def test_reranking_enabled_reorders_candidates(mocker: MockerFixture) -> None:
|
||||
def fake_post(_url: str, *, json: dict[str, object], timeout: float) -> httpx.Response:
|
||||
assert timeout == 30
|
||||
assert json == {
|
||||
@@ -77,16 +84,16 @@ def test_reranking_enabled_reorders_candidates(mocker: MockerFixture) -> None:
|
||||
}
|
||||
)
|
||||
|
||||
mocker.patch.object(httpx, "post", side_effect=fake_post)
|
||||
client = make_async_client(mocker, fake_post)
|
||||
|
||||
results = rerank_chunks("query", candidates(), RerankConfig())
|
||||
results = await rerank_chunks(client, "query", candidates(), RerankConfig())
|
||||
|
||||
assert [result.chunk_id for result in results] == [2, 1, 3]
|
||||
assert [round(result.score, 3) for result in results] == [0.78, 0.37, 0.28]
|
||||
assert [result.rerank_score for result in results] == [0.9, 0.1, 0.4]
|
||||
|
||||
|
||||
def test_reranking_cannot_ignore_hybrid_score(mocker: MockerFixture) -> None:
|
||||
async def test_reranking_cannot_ignore_hybrid_score(mocker: MockerFixture) -> None:
|
||||
candidates = [
|
||||
SearchResult(chunk_id=1, text="strong hybrid", source_title="A", score=1.0),
|
||||
SearchResult(chunk_id=2, text="weak hybrid", source_title="B", score=0.1),
|
||||
@@ -102,9 +109,9 @@ def test_reranking_cannot_ignore_hybrid_score(mocker: MockerFixture) -> None:
|
||||
}
|
||||
)
|
||||
|
||||
mocker.patch.object(httpx, "post", side_effect=fake_post)
|
||||
client = make_async_client(mocker, fake_post)
|
||||
|
||||
results = rerank_chunks("query", candidates, RerankConfig())
|
||||
results = await rerank_chunks(client, "query", candidates, RerankConfig())
|
||||
|
||||
assert [result.chunk_id for result in results] == [1, 2]
|
||||
assert results[0].score == pytest.approx(0.79)
|
||||
@@ -112,8 +119,9 @@ def test_reranking_cannot_ignore_hybrid_score(mocker: MockerFixture) -> None:
|
||||
assert results[1].rerank_score == 1.0
|
||||
|
||||
|
||||
def test_vllm_rerank_timeout_raises(mocker: MockerFixture) -> None:
|
||||
async def test_vllm_rerank_timeout_raises(mocker: MockerFixture) -> None:
|
||||
def fake_rerank_chunks(
|
||||
_client: httpx.AsyncClient,
|
||||
_query: str,
|
||||
_candidates: list[SearchResult],
|
||||
_config: RerankConfig,
|
||||
@@ -125,21 +133,21 @@ def test_vllm_rerank_timeout_raises(mocker: MockerFixture) -> None:
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=True), top_k=2)
|
||||
|
||||
with pytest.raises(httpx.TimeoutException, match="timeout"):
|
||||
apply_rerank("query", candidates(), config)
|
||||
await apply_rerank(mocker.Mock(), "query", candidates(), config)
|
||||
|
||||
|
||||
def test_malformed_vllm_rerank_json_does_not_crash_search(mocker: MockerFixture) -> None:
|
||||
async def test_malformed_vllm_rerank_json_does_not_crash_search(mocker: MockerFixture) -> None:
|
||||
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||
return rerank_response(content=b"not-json")
|
||||
|
||||
mocker.patch.object(httpx, "post", side_effect=fake_post)
|
||||
client = make_async_client(mocker, fake_post)
|
||||
|
||||
results = rerank_chunks("query", candidates()[:1], RerankConfig())
|
||||
results = await rerank_chunks(client, "query", candidates()[:1], RerankConfig())
|
||||
|
||||
assert results[0].score == 0.3
|
||||
|
||||
|
||||
def test_vllm_rerank_scores_are_clamped(mocker: MockerFixture) -> None:
|
||||
async def test_vllm_rerank_scores_are_clamped(mocker: MockerFixture) -> None:
|
||||
def fake_post(_url: str, **_kwargs: object) -> httpx.Response:
|
||||
return rerank_response(
|
||||
{
|
||||
@@ -150,8 +158,8 @@ def test_vllm_rerank_scores_are_clamped(mocker: MockerFixture) -> None:
|
||||
}
|
||||
)
|
||||
|
||||
mocker.patch.object(httpx, "post", side_effect=fake_post)
|
||||
client = make_async_client(mocker, fake_post)
|
||||
|
||||
results = rerank_chunks("query", candidates()[:2], RerankConfig())
|
||||
results = await rerank_chunks(client, "query", candidates()[:2], RerankConfig())
|
||||
|
||||
assert {result.chunk_id: result.rerank_score for result in results} == {1: 0.0, 2: 1.0}
|
||||
|
||||
+370
-21
@@ -2,32 +2,49 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from compression import zstd
|
||||
from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi import BackgroundTasks
|
||||
from fastapi.testclient import TestClient
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
|
||||
from sqlalchemy.pool import StaticPool
|
||||
|
||||
from python.ebook_search.api.bm25_tasks import refresh_bm25_for_engine
|
||||
from python.ebook_search.api.judge_tasks import (
|
||||
is_judging_book,
|
||||
judge_book_phrases_for_app,
|
||||
pop_book_judgment_outcome,
|
||||
start_book_phrase_judgment,
|
||||
)
|
||||
from python.ebook_search.api.main import create_app
|
||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||
from python.ebook_search.embeddings import EmbeddingModelStats
|
||||
from python.ebook_search.protected_phrases.models import (
|
||||
CorpusPhraseStats,
|
||||
PhraseCandidateGenerationResult,
|
||||
PhraseJudgmentBackfillResult,
|
||||
)
|
||||
from python.ebook_search.search import SearchResponse, SearchResult
|
||||
from python.ebook_search.timing import RuntimeStep
|
||||
from python.orm.richie import EbookSource, RichieBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from pytest_mock import MockerFixture
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
|
||||
def patch_app_runtime(mocker: MockerFixture):
|
||||
"""Patch app startup dependencies used by UI route tests."""
|
||||
mocker.patch("python.ebook_search.api.main.get_postgres_engine", side_effect=fake_get_postgres_engine)
|
||||
mocker.patch("python.ebook_search.api.main.get_async_postgres_engine", side_effect=fake_get_postgres_engine)
|
||||
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
|
||||
|
||||
|
||||
def fake_get_postgres_engine(**_kwargs):
|
||||
"""Return an in-memory engine for route tests."""
|
||||
return create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||
return create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
|
||||
|
||||
def test_search_page_uses_zstd_when_requested(mocker: MockerFixture) -> None:
|
||||
@@ -43,36 +60,68 @@ def test_search_page_uses_zstd_when_requested(mocker: MockerFixture) -> None:
|
||||
assert b"EPUB Search" in zstd.decompress(response.content)
|
||||
|
||||
|
||||
def test_ui_form_passes_rerank_flag_to_search_handler(mocker: MockerFixture) -> None:
|
||||
def test_ui_form_passes_search_toggles_to_search_handler(mocker: MockerFixture) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_search_ebooks(_engine, query, config, *, rerank=False):
|
||||
def fake_search_ebooks(_engine, _client, query, config, *, rerank=False, phrase_matching=False):
|
||||
captured["query"] = query
|
||||
captured["rerank"] = rerank
|
||||
captured["phrase_matching"] = phrase_matching
|
||||
captured["config"] = config
|
||||
return SearchResponse(query=query, results=[], rank_label="Hybrid + rerank")
|
||||
|
||||
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.search.answer_query",
|
||||
side_effect=lambda _query, _results, _config: "answer",
|
||||
side_effect=lambda _client, _query, _results, _config: "answer",
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.post("/search", data={"query": "where is the quote?", "rerank": "true"})
|
||||
response = client.post(
|
||||
"/search",
|
||||
data={"query": "where is the quote?", "rerank": "true", "phrase_matching": "true"},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "Hybrid + rerank" in response.text
|
||||
assert captured["query"] == "where is the quote?"
|
||||
assert captured["rerank"] is True
|
||||
assert captured["phrase_matching"] is True
|
||||
|
||||
|
||||
def test_ui_form_can_disable_phrase_matching(mocker: MockerFixture) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
captured["query"] = query
|
||||
captured["phrase_matching"] = phrase_matching
|
||||
return SearchResponse(query=query, results=[], rank_label="Hybrid")
|
||||
|
||||
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.search.answer_query",
|
||||
side_effect=lambda _client, _query, _results, _config: "answer",
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=12, answer_enabled=True)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.post("/search", data={"query": "where is the quote?"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured["query"] == "where is the quote?"
|
||||
assert captured["phrase_matching"] is False
|
||||
|
||||
|
||||
def test_ui_search_failure_returns_visible_error(mocker: MockerFixture) -> None:
|
||||
def fake_search_ebooks(_engine, _query, _config, *, rerank=False):
|
||||
def fake_search_ebooks(_engine, _client, _query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
del phrase_matching
|
||||
msg = "search exploded"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
@@ -89,11 +138,12 @@ def test_ui_search_failure_returns_visible_error(mocker: MockerFixture) -> None:
|
||||
|
||||
|
||||
def test_ui_answer_failure_still_returns_sources(mocker: MockerFixture) -> None:
|
||||
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
del phrase_matching
|
||||
return SearchResponse(query=query, results=[], rank_label="Hybrid")
|
||||
|
||||
def fake_answer_query(_query, _results, _config):
|
||||
def fake_answer_query(_client, _query, _results, _config):
|
||||
msg = "answer exploded"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
@@ -113,11 +163,12 @@ def test_ui_answer_failure_still_returns_sources(mocker: MockerFixture) -> None:
|
||||
def test_ui_skips_answer_when_disabled(mocker: MockerFixture) -> None:
|
||||
called = False
|
||||
|
||||
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
del phrase_matching
|
||||
return SearchResponse(query=query, results=[], rank_label="Hybrid")
|
||||
|
||||
def fake_answer_query(_query, _results, _config):
|
||||
def fake_answer_query(_client, _query, _results, _config):
|
||||
nonlocal called
|
||||
called = True
|
||||
return "answer"
|
||||
@@ -138,8 +189,9 @@ def test_ui_skips_answer_when_disabled(mocker: MockerFixture) -> None:
|
||||
|
||||
|
||||
def test_ui_shows_component_scores(mocker: MockerFixture) -> None:
|
||||
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
del phrase_matching
|
||||
return SearchResponse(
|
||||
query=query,
|
||||
rank_label="Hybrid + rerank",
|
||||
@@ -160,7 +212,7 @@ def test_ui_shows_component_scores(mocker: MockerFixture) -> None:
|
||||
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.search.answer_query",
|
||||
side_effect=lambda _query, _results, _config: "answer",
|
||||
side_effect=lambda _client, _query, _results, _config: "answer",
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
@@ -176,9 +228,47 @@ def test_ui_shows_component_scores(mocker: MockerFixture) -> None:
|
||||
assert "RRF" in response.text
|
||||
|
||||
|
||||
def test_ui_shows_search_runtime_chart(mocker: MockerFixture) -> None:
|
||||
def fake_search_ebooks(_engine, query, _config, *, rerank=False):
|
||||
def test_ui_shows_matched_phrases_that_boosted_a_result(mocker: MockerFixture) -> None:
|
||||
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
del phrase_matching
|
||||
return SearchResponse(
|
||||
query=query,
|
||||
rank_label="Hybrid",
|
||||
results=[
|
||||
SearchResult(
|
||||
chunk_id=1,
|
||||
text="source text",
|
||||
source_title="Book",
|
||||
score=0.9,
|
||||
phrase_hit_count=3,
|
||||
matched_phrases=("lock in", "haden's syndrome"),
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.search.answer_query",
|
||||
side_effect=lambda _client, _query, _results, _config: "answer",
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False), answer_enabled=True)
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.post("/search", data={"query": "what is lock in?"})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "boosted by" in response.text
|
||||
assert "lock in" in response.text
|
||||
assert "haden's syndrome" in response.text
|
||||
|
||||
|
||||
def test_ui_shows_search_runtime_chart(mocker: MockerFixture) -> None:
|
||||
def fake_search_ebooks(_engine, _client, query, _config, *, rerank=False, phrase_matching=False):
|
||||
del rerank
|
||||
del phrase_matching
|
||||
return SearchResponse(
|
||||
query=query,
|
||||
rank_label="Hybrid",
|
||||
@@ -192,7 +282,7 @@ def test_ui_shows_search_runtime_chart(mocker: MockerFixture) -> None:
|
||||
mocker.patch("python.ebook_search.api.routes.search.search_ebooks", side_effect=fake_search_ebooks)
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.search.answer_query",
|
||||
side_effect=lambda _query, _results, _config: "answer",
|
||||
side_effect=lambda _client, _query, _results, _config: "answer",
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
@@ -214,7 +304,7 @@ def test_ui_embed_all_batches_until_complete(mocker: MockerFixture) -> None:
|
||||
counts = iter([32, 32, 5, 0])
|
||||
batch_sizes: list[int] = []
|
||||
|
||||
def fake_embed_missing_chunks(_session, config):
|
||||
def fake_embed_missing_chunks(_session, _client, config):
|
||||
batch_sizes.append(config.embedding_batch_size)
|
||||
return next(counts)
|
||||
|
||||
@@ -256,7 +346,7 @@ def test_ui_scan_schedules_bm25_refresh_after_database_change(mocker: MockerFixt
|
||||
assert scheduled is True
|
||||
|
||||
|
||||
def test_bm25_refresh_clears_loaded_corpus_cache(mocker: MockerFixture) -> None:
|
||||
async def test_bm25_refresh_clears_loaded_corpus_cache(mocker: MockerFixture) -> None:
|
||||
refreshed: list[object] = []
|
||||
cache_cleared = False
|
||||
|
||||
@@ -269,16 +359,146 @@ def test_bm25_refresh_clears_loaded_corpus_cache(mocker: MockerFixture) -> None:
|
||||
|
||||
mocker.patch("python.ebook_search.api.bm25_tasks.refresh_bm25_corpus", side_effect=fake_refresh_bm25_corpus)
|
||||
mocker.patch("python.ebook_search.api.bm25_tasks.load_bm25_corpus.cache_clear", side_effect=fake_cache_clear)
|
||||
engine = create_engine("sqlite+pysqlite:///:memory:", future=True)
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
|
||||
refresh_bm25_for_engine(engine, config)
|
||||
await refresh_bm25_for_engine(engine, config)
|
||||
|
||||
assert len(refreshed) == 1
|
||||
assert refreshed[0][1] == config
|
||||
assert cache_cleared is True
|
||||
|
||||
|
||||
def build_engine_with_book() -> AsyncEngine:
|
||||
"""Create a shareable in-memory async engine holding one indexed book."""
|
||||
engine = create_async_engine(
|
||||
"sqlite+aiosqlite:///:memory:",
|
||||
connect_args={"check_same_thread": False},
|
||||
poolclass=StaticPool,
|
||||
)
|
||||
|
||||
async def seed() -> None:
|
||||
async with engine.begin() as connection:
|
||||
await connection.run_sync(RichieBase.metadata.create_all)
|
||||
async with AsyncSession(engine) as session:
|
||||
session.add(
|
||||
EbookSource(
|
||||
title="Book",
|
||||
author="Author",
|
||||
language=None,
|
||||
publisher=None,
|
||||
identifier=None,
|
||||
file_path="/library/book.epub",
|
||||
file_sha256="a" * 64,
|
||||
file_mtime=datetime.now(tz=UTC),
|
||||
file_size=10,
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
asyncio.run(seed())
|
||||
return engine
|
||||
|
||||
|
||||
def test_ui_judge_phrases_redirects_and_judges_in_background(mocker: MockerFixture) -> None:
|
||||
mocker.patch("python.ebook_search.api.main.get_async_postgres_engine", return_value=build_engine_with_book())
|
||||
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
|
||||
judged_source_ids: list[list[int]] = []
|
||||
|
||||
def fake_judge(_engine: object, _config: object, *, source_ids: list[int]) -> PhraseJudgmentBackfillResult:
|
||||
judged_source_ids.append(source_ids)
|
||||
return PhraseJudgmentBackfillResult(
|
||||
books_seen=1,
|
||||
books_judged=1,
|
||||
books_failed=0,
|
||||
candidates_judged=3,
|
||||
protected_phrases=2,
|
||||
phrase_mentions=4,
|
||||
)
|
||||
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.judge_tasks.judge_candidate_phrases_for_books",
|
||||
side_effect=fake_judge,
|
||||
)
|
||||
app = create_app()
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.post("/books/1/judge-phrases", follow_redirects=False)
|
||||
detail_after = client.get("/books/1")
|
||||
detail_again = client.get("/books/1")
|
||||
|
||||
assert response.status_code == 303
|
||||
assert response.headers["location"] == "/books/1"
|
||||
assert judged_source_ids == [[1]]
|
||||
assert "Judged 3 candidates; 2 protected phrases promoted" in detail_after.text
|
||||
assert "Judged 3 candidates" not in detail_again.text
|
||||
|
||||
|
||||
def test_ui_book_detail_shows_judging_in_progress(mocker: MockerFixture) -> None:
|
||||
mocker.patch("python.ebook_search.api.main.get_async_postgres_engine", return_value=build_engine_with_book())
|
||||
mocker.patch("python.ebook_search.api.main.ensure_bm25_corpus", side_effect=lambda _session, _config: None)
|
||||
mocker.patch("python.ebook_search.api.routes.page.is_judging_book", return_value=True)
|
||||
app = create_app()
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/books/1")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "Judging candidate phrases in the background" in response.text
|
||||
assert "disabled" in response.text
|
||||
|
||||
|
||||
def test_book_phrase_judgment_rejects_duplicate_while_queued(mocker: MockerFixture) -> None:
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.judge_tasks.judge_candidate_phrases_for_books",
|
||||
return_value=PhraseJudgmentBackfillResult(
|
||||
books_seen=1,
|
||||
books_judged=1,
|
||||
books_failed=0,
|
||||
candidates_judged=3,
|
||||
protected_phrases=2,
|
||||
phrase_mentions=4,
|
||||
),
|
||||
)
|
||||
app = create_app()
|
||||
app.state.engine = None
|
||||
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
background_tasks = BackgroundTasks()
|
||||
|
||||
assert start_book_phrase_judgment(app, background_tasks, 1) is True
|
||||
assert is_judging_book(app, 1) is True
|
||||
assert start_book_phrase_judgment(app, background_tasks, 1) is False
|
||||
assert len(background_tasks.tasks) == 1
|
||||
|
||||
asyncio.run(judge_book_phrases_for_app(app, 1))
|
||||
|
||||
assert is_judging_book(app, 1) is False
|
||||
assert pop_book_judgment_outcome(app, 1) == "Judged 3 candidates; 2 protected phrases promoted"
|
||||
assert pop_book_judgment_outcome(app, 1) is None
|
||||
assert start_book_phrase_judgment(app, background_tasks, 1) is True
|
||||
|
||||
|
||||
def test_book_phrase_judgment_records_failure_outcome(mocker: MockerFixture) -> None:
|
||||
def fake_judge(_engine: object, _config: object, *, source_ids: list[int]) -> PhraseJudgmentBackfillResult:
|
||||
del source_ids
|
||||
message = "llm judge unavailable"
|
||||
raise RuntimeError(message)
|
||||
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.judge_tasks.judge_candidate_phrases_for_books",
|
||||
side_effect=fake_judge,
|
||||
)
|
||||
app = create_app()
|
||||
app.state.engine = None
|
||||
app.state.config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
|
||||
start_book_phrase_judgment(app, BackgroundTasks(), 7)
|
||||
asyncio.run(judge_book_phrases_for_app(app, 7))
|
||||
|
||||
assert is_judging_book(app, 7) is False
|
||||
assert pop_book_judgment_outcome(app, 7) == "Judging failed; see server logs for details"
|
||||
|
||||
|
||||
def test_admin_page_shows_embedding_counts_by_model(mocker: MockerFixture) -> None:
|
||||
def fake_embedding_model_stats(_session):
|
||||
return [
|
||||
@@ -297,6 +517,10 @@ def test_admin_page_shows_embedding_counts_by_model(mocker: MockerFixture) -> No
|
||||
]
|
||||
|
||||
mocker.patch("python.ebook_search.api.routes.admin.embedding_model_stats", side_effect=fake_embedding_model_stats)
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.admin.corpus_phrase_stats",
|
||||
return_value=fake_corpus_phrase_stats(),
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
|
||||
@@ -310,3 +534,128 @@ def test_admin_page_shows_embedding_counts_by_model(mocker: MockerFixture) -> No
|
||||
assert "24" in response.text
|
||||
assert "qwen3-embedding-4b" in response.text
|
||||
assert "2560" in response.text
|
||||
|
||||
|
||||
def fake_corpus_phrase_stats() -> CorpusPhraseStats:
|
||||
"""Build distinctive corpus phrase stats for admin page assertions."""
|
||||
return CorpusPhraseStats(
|
||||
total_books=17,
|
||||
books_with_candidates=13,
|
||||
books_fully_judged=11,
|
||||
candidate_phrases=901,
|
||||
judged_candidates=703,
|
||||
unjudged_candidates=198,
|
||||
protected_phrases=157,
|
||||
)
|
||||
|
||||
|
||||
def test_admin_page_shows_protected_phrase_stats(mocker: MockerFixture) -> None:
|
||||
mocker.patch("python.ebook_search.api.routes.admin.embedding_model_stats", return_value=[])
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.admin.corpus_phrase_stats",
|
||||
return_value=fake_corpus_phrase_stats(),
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/admin")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "Protected phrases" in response.text
|
||||
for value in ("17", "13", "11", "901", "703", "198", "157"):
|
||||
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
|
||||
return PhraseCandidateGenerationResult(books_seen=5, books_built=5, candidate_phrases=99)
|
||||
|
||||
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-all")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured["only_missing"] is False
|
||||
assert "5 of 5 books" in response.text
|
||||
|
||||
|
||||
def test_ui_judge_missing_phrases_judges_only_pending_books(mocker: MockerFixture) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
async def fake_judge(_engine, _config, *, source_ids=None):
|
||||
captured["source_ids"] = source_ids
|
||||
return PhraseJudgmentBackfillResult(
|
||||
books_seen=2,
|
||||
books_judged=2,
|
||||
books_failed=0,
|
||||
candidates_judged=10,
|
||||
protected_phrases=4,
|
||||
phrase_mentions=9,
|
||||
)
|
||||
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.admin.judge_candidate_phrases_for_books",
|
||||
side_effect=fake_judge,
|
||||
)
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.admin.book_ids_pending_first_judgment",
|
||||
return_value=[3, 5],
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.post("/admin/phrases/judge-missing")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert captured["source_ids"] == [3, 5]
|
||||
assert "4 protected phrases" in response.text
|
||||
|
||||
|
||||
def test_ui_judge_missing_phrases_reports_when_nothing_is_pending(mocker: MockerFixture) -> None:
|
||||
judge = mocker.patch("python.ebook_search.api.routes.admin.judge_candidate_phrases_for_books")
|
||||
mocker.patch(
|
||||
"python.ebook_search.api.routes.admin.book_ids_pending_first_judgment",
|
||||
return_value=[],
|
||||
)
|
||||
patch_app_runtime(mocker)
|
||||
app = create_app()
|
||||
|
||||
with TestClient(app) as client:
|
||||
response = client.post("/admin/phrases/judge-missing")
|
||||
|
||||
assert response.status_code == 200
|
||||
assert "have been judged" in response.text
|
||||
judge.assert_not_called()
|
||||
|
||||
Reference in New Issue
Block a user