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

refactor(ebook-search): simplify search and phrase matching

Browse Source
This commit is contained in:
Richie
2026-07-24 11:38:51 -04:00
parent 0028237579
commit e010756e09
16 changed files with 490 additions and 557 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/ebook_search/test_core.py
+8 -7
View File
@@ -25,7 +25,7 @@ from python.ebook_search.bm25_corpus import (
score_bm25_corpus,
write_bm25_corpus,
)
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config, normalize_embedding_model
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config
from python.ebook_search.embeddings import MODEL_DIMENSIONS, ensure_embedding_models
from python.ebook_search.ingest import chunk_text, find_existing_source
from python.ebook_search.search import (
@@ -452,24 +452,25 @@ def test_1024_embedding_table_has_cosine_hnsw_index() -> None:
def test_embedding_model_aliases_normalize_to_provider_names(mocker: MockerFixture) -> None:
mocker.patch.dict(environ, {}, clear=False)
environ.pop("EBOOK_SEARCH_EMBEDDING_MODEL", None)
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
assert load_config().embedding_model == "qwen3-embedding-0.6b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-0.6b"
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
assert load_config().embedding_model == "qwen3-embedding-0.6b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen3-Embedding-0.6B"
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
assert load_config().embedding_model == "qwen3-embedding-0.6b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen/Qwen3-Embedding-4B"
assert normalize_embedding_model() == "qwen3-embedding-4b"
assert load_config().embedding_model == "qwen3-embedding-4b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding:8b"
assert normalize_embedding_model() == "qwen3-embedding-8b"
assert load_config().embedding_model == "qwen3-embedding-8b"
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-8b"
assert normalize_embedding_model() == "qwen3-embedding-8b"
assert load_config().embedding_model == "qwen3-embedding-8b"
def test_answer_generation_is_enabled_by_default(mocker: MockerFixture) -> None:
tests/ebook_search/test_http.py
+39
View File
@@ -10,6 +10,7 @@ import pytest
from python.ebook_search.answer import answer_query
from python.ebook_search.config import EbookSearchConfig, RerankConfig
from python.ebook_search.embeddings import embed_texts
from python.ebook_search.llm_interface import check_chat_endpoint, check_embedding_endpoint
from python.ebook_search.search import SearchResult
if TYPE_CHECKING:
@@ -23,6 +24,44 @@ def make_async_client(mocker: MockerFixture, fake_post) -> httpx.AsyncClient:
return client
async def test_model_endpoint_checks_share_http_probe(mocker: MockerFixture) -> None:
client = mocker.MagicMock(spec=httpx.AsyncClient)
response = mocker.MagicMock(spec=httpx.Response)
client.get = mocker.AsyncMock(return_value=response)
config = EbookSearchConfig(
rerank=RerankConfig(enabled=False),
embedding_base_url="https://embedding.example/v1/",
vllm_base_url="https://chat.example/v1/",
vllm_api_key="secret",
)
assert await check_embedding_endpoint(client, config, timeout_seconds=2.0)
assert await check_chat_endpoint(client, config, timeout_seconds=3.0)
assert client.get.await_args_list == [
mocker.call("https://embedding.example/v1/models", headers={}, timeout=2.0),
mocker.call(
"https://chat.example/v1/models",
headers={"Authorization": "Bearer secret"},
timeout=3.0,
),
]
assert response.raise_for_status.call_count == 2
async def test_model_endpoint_checks_report_http_failures(mocker: MockerFixture) -> None:
client = mocker.MagicMock(spec=httpx.AsyncClient)
client.get = mocker.AsyncMock(
side_effect=[
httpx.ConnectError("embedding offline"),
httpx.ConnectError("chat offline"),
]
)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
assert not await check_embedding_endpoint(client, config)
assert not await check_chat_endpoint(client, config)
async def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None:
captured: dict[str, object] = {}
tests/ebook_search/test_protected_phrases.py
+89 -23
View File
@@ -7,7 +7,7 @@ from datetime import UTC, datetime
from typing import TYPE_CHECKING
import pytest
from sqlalchemy import select
from sqlalchemy import event, select
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
from python.ebook_search.config import EbookSearchConfig
@@ -33,15 +33,12 @@ 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,
PhraseMatch,
)
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
@@ -187,26 +184,89 @@ def test_score_candidate_weights_metadata_source(config: EbookSearchConfig) -> N
assert score_candidate(metadata_only, config) == 2.0 + 0.5
async def test_detect_protected_phrases_hydrates_alias_matches(
async def test_load_phrase_lookup_loads_phrases_and_aliases_with_one_query(
engine: AsyncEngine,
session: AsyncSession,
) -> None:
"""One joined query should load phrases with zero, one, or many aliases."""
source = await add_source(session)
aliased_phrase = await add_phrase(session, source.id, phrase_text="lock in", phrase_norm="lock in")
plain_phrase = await add_phrase(session, source.id, phrase_text="mage king", phrase_norm="mage king")
session.add_all(
[
EbookPhraseAlias(phrase_id=aliased_phrase.id, alias_text="locked in", alias_norm="locked in"),
EbookPhraseAlias(
phrase_id=aliased_phrase.id,
alias_text="locked completely in",
alias_norm="locked completely in",
),
]
)
await session.commit()
statements: list[str] = []
def record_statement(*args: object) -> None:
statements.append(str(args[2]))
event.listen(engine.sync_engine, "before_cursor_execute", record_statement)
try:
lookup = await load_phrase_lookup(session, EbookSearchConfig(phrase_max_tokens=2), book_id=source.id)
finally:
event.remove(engine.sync_engine, "before_cursor_execute", record_statement)
assert lookup.phrase_ids_by_norm == {
"lock in": (aliased_phrase.id,),
"locked completely in": (aliased_phrase.id,),
"locked in": (aliased_phrase.id,),
"mage king": (plain_phrase.id,),
}
assert lookup.phrases_by_id == {aliased_phrase.id: aliased_phrase, plain_phrase.id: plain_phrase}
assert lookup.max_tokens == 3
assert len(statements) == 1
async def test_detect_protected_phrases_queries_canonical_and_alias_matches_once(
engine: AsyncEngine,
session: AsyncSession,
config: EbookSearchConfig,
) -> None:
"""Query detection should use RAM aliases and hydrate phrase metadata from the DB."""
"""Query detection should hydrate repeated canonical and alias matches with one statement."""
source = await add_source(session)
phrase = await add_phrase(session, source.id, phrase_text="lock in", phrase_norm="lock in")
await add_phrase(session, source.id, phrase_text="mage king", phrase_norm="mage king")
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)
statements: list[str] = []
assert [(match.phrase_text, match.canonical_id, match.phrase_type) for match in matches] == [
("lock in", "condition:lock_in", "fictional_condition")
def record_statement(*args: object) -> None:
statements.append(str(args[2]))
event.listen(engine.sync_engine, "before_cursor_execute", record_statement)
try:
matches = await detect_protected_phrases_for_query(
session,
"lock-in then locked-in and lock-in",
config,
)
empty_matches = await detect_protected_phrases_for_query(session, "", config)
finally:
event.remove(engine.sync_engine, "before_cursor_execute", record_statement)
assert [(match.phrase_text, match.matched_norm, match.start_token, match.end_token) for match in matches] == [
("lock in", "lock in", 0, 2),
("lock in", "locked in", 3, 5),
("lock in", "lock in", 6, 8),
]
assert empty_matches == []
assert len(statements) == 1
assert " UNION " in statements[0]
def test_resolve_overlaps_keeps_independent_nested_phrases() -> None:
"""Overlap resolution should keep useful nested concepts when metadata permits it."""
child = hydrated_match(
child = phrase_match(
phrase_id=1,
phrase_text="mage king",
canonical_id="title:mage_king",
@@ -214,7 +274,7 @@ def test_resolve_overlaps_keeps_independent_nested_phrases() -> None:
end_token=5,
allow_nested=True,
)
parent = hydrated_match(
parent = phrase_match(
phrase_id=2,
phrase_text="mage king of mars",
canonical_id="entity:mage_king_of_mars",
@@ -231,7 +291,7 @@ def test_resolve_overlaps_keeps_independent_nested_phrases() -> None:
def test_resolve_overlaps_suppresses_weaker_same_canonical_match() -> None:
"""Same-canonical overlaps should keep the stronger evidence."""
weak = hydrated_match(
weak = phrase_match(
phrase_id=1,
phrase_text="lock",
canonical_id="condition:lock_in",
@@ -239,7 +299,7 @@ def test_resolve_overlaps_suppresses_weaker_same_canonical_match() -> None:
end_token=3,
importance=0.2,
)
strong = hydrated_match(
strong = phrase_match(
phrase_id=2,
phrase_text="lock in",
canonical_id="condition:lock_in",
@@ -274,19 +334,25 @@ async def test_index_chunk_phrase_mentions_uses_normalized_window_lookup(
await session.commit()
lookup = await load_phrase_lookup(session, config, book_id=source.id)
count = await index_chunk_phrase_mentions(session, chunk, lookup=lookup)
statements: list[str] = []
def record_statement(*args: object) -> None:
statements.append(str(args[2]))
event.listen(session.bind.sync_engine, "before_cursor_execute", record_statement)
try:
count = index_chunk_phrase_mentions(session, chunk, lookup=lookup)
finally:
event.remove(session.bind.sync_engine, "before_cursor_execute", record_statement)
await session.commit()
assert statements == []
mention = await session.scalar(select(EbookChunkPhraseMention))
assert count == 1
assert mention is not None
assert mention.chunk_id == chunk.id
assert mention.phrase_id == phrase.id
assert chunk.text[mention.start_char : mention.end_char] == "lock-in"
assert await phrase_hit_counts_for_chunks(session, chunk_ids=[chunk.id], phrase_ids=[phrase.id]) == {chunk.id: 1}
assert await phrase_hits_for_chunks(session, chunk_ids=[chunk.id], phrase_ids=[phrase.id]) == {
chunk.id: (ChunkPhraseHit(phrase_id=phrase.id, phrase_text="lock in", mention_count=1),)
}
@pytest.mark.usefixtures("worker_pool")
@@ -1060,7 +1126,7 @@ async def add_phrase(
return phrase
def hydrated_match(
def phrase_match(
*,
phrase_id: int,
phrase_text: str,
@@ -1070,9 +1136,9 @@ def hydrated_match(
importance: float = 0.8,
allow_nested: bool = False,
suppress_children: bool = True,
) -> HydratedPhraseMatch:
"""Build a hydrated match for overlap tests."""
return HydratedPhraseMatch(
) -> PhraseMatch:
"""Build a metadata-backed phrase match for overlap tests."""
return PhraseMatch(
phrase_id=phrase_id,
matched_norm=phrase_text,
phrase_text=phrase_text,
tests/ebook_search/test_rag_pipeline.py
+4 -3
View File
@@ -62,8 +62,9 @@ async def test_search_ebooks_runs_phrase_detection_in_parallel_with_retrieval(mo
assert await asyncio.to_thread(phrase_started.wait, 2)
return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]
async def fake_query_phrase_matches(_engine, _query, _config):
async def fake_query_phrase_matches(_engine, _query, _config, *, phrase_matching):
"""Record that phrase detection started and return no matches."""
assert phrase_matching is True
phrase_started.set()
return []
@@ -90,7 +91,7 @@ async def test_search_ebooks_skips_phrase_matching_when_disabled(mocker: MockerF
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")
detect_mock = mocker.patch("python.ebook_search.search.detect_protected_phrases_for_query")
boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
@@ -115,7 +116,7 @@ async def test_search_ebooks_ignores_phrase_matching_when_config_disabled(mocker
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")
detect_mock = mocker.patch("python.ebook_search.search.detect_protected_phrases_for_query")
boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), phrase_matching_enabled=False)