diff --git a/tests/ebook_search/test_core.py b/tests/ebook_search/test_core.py index 64a517e..61ac601 100644 --- a/tests/ebook_search/test_core.py +++ b/tests/ebook_search/test_core.py @@ -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 diff --git a/tests/ebook_search/test_guardrails.py b/tests/ebook_search/test_guardrails.py index 03cddfd..4436160 100644 --- a/tests/ebook_search/test_guardrails.py +++ b/tests/ebook_search/test_guardrails.py @@ -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() diff --git a/tests/ebook_search/test_health.py b/tests/ebook_search/test_health.py index be0ef35..2d8b5fb 100644 --- a/tests/ebook_search/test_health.py +++ b/tests/ebook_search/test_health.py @@ -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) diff --git a/tests/ebook_search/test_http.py b/tests/ebook_search/test_http.py index 2d2adec..f8f6266 100644 --- a/tests/ebook_search/test_http.py +++ b/tests/ebook_search/test_http.py @@ -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) diff --git a/tests/ebook_search/test_protected_phrases.py b/tests/ebook_search/test_protected_phrases.py new file mode 100644 index 0000000..8444142 --- /dev/null +++ b/tests/ebook_search/test_protected_phrases.py @@ -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, + ) diff --git a/tests/ebook_search/test_rag_pipeline.py b/tests/ebook_search/test_rag_pipeline.py index cc2a2b5..598d79c 100644 --- a/tests/ebook_search/test_rag_pipeline.py +++ b/tests/ebook_search/test_rag_pipeline.py @@ -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() diff --git a/tests/ebook_search/test_rerank.py b/tests/ebook_search/test_rerank.py index b421832..0fcf907 100644 --- a/tests/ebook_search/test_rerank.py +++ b/tests/ebook_search/test_rerank.py @@ -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} diff --git a/tests/ebook_search/test_ui.py b/tests/ebook_search/test_ui.py index 8db5094..89db115 100644 --- a/tests/ebook_search/test_ui.py +++ b/tests/ebook_search/test_ui.py @@ -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()