From 70bf8627a26d6493892195f463014fbd7913d55d Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Thu, 2 Jul 2026 16:30:14 -0400 Subject: [PATCH] feat(ebook): implement phrase matching functionality and UI enhancements --- python/ebook_search/api/main.py | 3 +- python/ebook_search/api/routes/search.py | 9 +- python/ebook_search/api/static/style.css | 6 + python/ebook_search/api/templates/search.html | 19 ++- python/ebook_search/bm25_corpus.py | 1 + python/ebook_search/config.py | 1 + python/ebook_search/ingest.py | 5 +- python/ebook_search/protected_phrases/lib.py | 3 +- python/ebook_search/search.py | 117 +++++++++++++++++- 9 files changed, 151 insertions(+), 13 deletions(-) diff --git a/python/ebook_search/api/main.py b/python/ebook_search/api/main.py index 286e8a7..87e3316 100644 --- a/python/ebook_search/api/main.py +++ b/python/ebook_search/api/main.py @@ -36,12 +36,13 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: app.state.config = config logger.info( "ebook_search_config_loaded top_k=%s embedding_model=%s embedding_base_url=%s vllm_base_url=%s " - "rerank_enabled=%s answer_enabled=%s library_paths=%s", + "rerank_enabled=%s phrase_matching_enabled=%s answer_enabled=%s library_paths=%s", config.top_k, config.embedding_model, config.embedding_base_url, config.vllm_base_url, config.rerank.enabled, + config.phrase_matching_enabled, config.answer_enabled, len(config.library_paths), ) diff --git a/python/ebook_search/api/routes/search.py b/python/ebook_search/api/routes/search.py index c2aa52a..6c7f8a6 100644 --- a/python/ebook_search/api/routes/search.py +++ b/python/ebook_search/api/routes/search.py @@ -80,10 +80,17 @@ def search( engine: AppEngine, query: Annotated[str, Form()], rerank: Annotated[str | None, Form()] = None, + phrase_matching: Annotated[str | None, Form()] = None, ) -> HTMLResponse: """Run a search and render HTMX results.""" try: - response = search_ebooks(engine, query, config, rerank=rerank == "true") + response = search_ebooks( + engine, + query, + config, + rerank=rerank == "true", + phrase_matching=phrase_matching == "true", + ) except Exception as error: logger.exception("ebook_search_request_failed") return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) diff --git a/python/ebook_search/api/static/style.css b/python/ebook_search/api/static/style.css index 0a74b4e..4644985 100644 --- a/python/ebook_search/api/static/style.css +++ b/python/ebook_search/api/static/style.css @@ -181,6 +181,12 @@ textarea:focus { flex-wrap: wrap; } +.search-toggles { + display: flex; + flex-wrap: wrap; + gap: 14px; +} + button { padding: 10px 20px; font: inherit; diff --git a/python/ebook_search/api/templates/search.html b/python/ebook_search/api/templates/search.html index 9ae3a43..c86113f 100644 --- a/python/ebook_search/api/templates/search.html +++ b/python/ebook_search/api/templates/search.html @@ -9,10 +9,21 @@
- +
+ + +
diff --git a/python/ebook_search/bm25_corpus.py b/python/ebook_search/bm25_corpus.py index 3a1752b..0450d64 100644 --- a/python/ebook_search/bm25_corpus.py +++ b/python/ebook_search/bm25_corpus.py @@ -174,6 +174,7 @@ def fetch_bm25_corpus_records(session: Session) -> tuple[list[dict[str, object]] select( EbookChunk.id.label("chunk_id"), EbookChunk.text.label("text"), + EbookSource.id.label("source_id"), EbookSource.title.label("source_title"), EbookSource.author.label("source_author"), EbookChapter.title.label("chapter_title"), diff --git a/python/ebook_search/config.py b/python/ebook_search/config.py index b2a7e09..173f07b 100644 --- a/python/ebook_search/config.py +++ b/python/ebook_search/config.py @@ -89,6 +89,7 @@ class EbookSearchConfig(BaseSettings): protected_phrase_max_candidates_per_book: int = 5000 protected_phrase_llm_candidates_per_book: int = 500 protected_phrase_confidence_threshold: float = 0.80 + phrase_matching_enabled: bool = True phrase_hit_boost: float = 0.25 phrase_min_tokens: int = 2 phrase_max_tokens: int = 5 diff --git a/python/ebook_search/ingest.py b/python/ebook_search/ingest.py index 516735d..4111f69 100644 --- a/python/ebook_search/ingest.py +++ b/python/ebook_search/ingest.py @@ -13,6 +13,7 @@ import tiktoken from sqlalchemy import or_, select from python.ebook_search.epub_parse import parse_epub +from python.ebook_search.protected_phrases.lib import index_chunk_phrase_mentions_for_book from python.orm.richie import EbookChapter, EbookChunk, EbookSource logger = logging.getLogger(__name__) @@ -138,12 +139,14 @@ def ingest_file(session: Session, path: Path, config: EbookSearchConfig) -> bool chunk_index = add_chapter_chunks(session, source, chapter, parsed_chapter, chunk_index, config) session.commit() + mention_count = index_chunk_phrase_mentions_for_book(session, source.id, config) logger.info( - "ebook_ingest_file_complete source_id=%s path=%s chapters=%s chunks=%s", + "ebook_ingest_file_complete source_id=%s path=%s chapters=%s chunks=%s phrase_mentions=%s", source.id, resolved_path, len(parsed.chapters), chunk_index, + mention_count, ) except Exception: logger.exception(f"ebook_ingest_file_error path={path}") diff --git a/python/ebook_search/protected_phrases/lib.py b/python/ebook_search/protected_phrases/lib.py index aec90b9..08d40bf 100644 --- a/python/ebook_search/protected_phrases/lib.py +++ b/python/ebook_search/protected_phrases/lib.py @@ -1249,8 +1249,7 @@ def judge_candidate_phrases_for_books( ) continue logger.info( - "ebook_candidate_phrase_judgment_book_start source_id=%s book_number=%s books_seen=%s " - "title=%r unjudged=%s", + "ebook_candidate_phrase_judgment_book_start source_id=%s book_number=%s books_seen=%s title=%r unjudged=%s", source.id, book_number, books_seen, diff --git a/python/ebook_search/search.py b/python/ebook_search/search.py index 6e3e8d1..b266c2a 100644 --- a/python/ebook_search/search.py +++ b/python/ebook_search/search.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING from pgvector.sqlalchemy import Vector from sqlalchemy import literal, select +from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session from python.ebook_search.bm25_corpus import ( @@ -19,6 +20,11 @@ from python.ebook_search.bm25_corpus import ( score_bm25_corpus, ) from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table +from python.ebook_search.protected_phrases.lib import ( + HydratedPhraseMatch, + detect_protected_phrases_for_query, + phrase_hit_counts_for_chunks, +) from python.ebook_search.rerank import rerank_chunks from python.ebook_search.timing import RuntimeStep, timed_result from python.orm.richie import ( @@ -29,7 +35,7 @@ from python.orm.richie import ( ) if TYPE_CHECKING: - from collections.abc import Mapping + from collections.abc import Mapping, Sequence from sqlalchemy.engine import Engine @@ -45,11 +51,13 @@ class SearchResult: chunk_id: int text: str source_title: str + source_id: int | None = None score: float = 0.0 vector_score: float | None = None bm25_score: float | None = None fused_score: float | None = None rerank_score: float | None = None + phrase_hit_count: int = 0 source_author: str | None = None chapter_title: str | None = None page_label: str | None = None @@ -64,6 +72,7 @@ class SearchResponse: results: list[SearchResult] rank_label: str timings: tuple[RuntimeStep, ...] = () + phrase_matches: tuple[HydratedPhraseMatch, ...] = () @property def total_runtime_ms(self) -> float: @@ -86,14 +95,26 @@ def search_ebooks( config: EbookSearchConfig, *, rerank: bool = False, + phrase_matching: bool | None = None, ) -> SearchResponse: """Run hybrid vector/BM25 search and optional reranking.""" if not query.strip(): logger.info("ebook_search_empty_query") return SearchResponse(query=query, results=[], rank_label="Hybrid") - logger.info("ebook_search_start query_length=%s rerank=%s", len(query), rerank) + phrase_matching_enabled = config.phrase_matching_enabled if phrase_matching is None else phrase_matching + logger.info( + "ebook_search_start query_length=%s rerank=%s phrase_matching=%s", + len(query), + rerank, + phrase_matching_enabled, + ) timings: list[RuntimeStep] = [] + if phrase_matching_enabled: + phrase_matches, timing = timed_result("Protected phrase detection", query_phrase_matches, engine, query, config) + else: + phrase_matches, timing = timed_result("Protected phrase detection skipped", skip_phrase_matches) + timings.append(timing) retrieval, timing = timed_result( "Hybrid retrieval", parallel_retrieval, @@ -111,18 +132,32 @@ def search_ebooks( rank_constant=config.rrf_rank_constant, ) timings.append(timing) + if phrase_matching_enabled: + fused, timing = timed_result( + "Phrase mention boost", + apply_phrase_mention_boosts, + engine, + fused, + phrase_matches, + config.phrase_hit_boost, + ) + else: + fused, timing = timed_result("Phrase mention boost skipped", skip_phrase_mention_boosts, fused) + timings.append(timing) if config.rerank.enabled and rerank: response, timing = timed_result("Rerank", apply_rerank, query, fused, config) else: response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config) timings.append(timing) - response = replace(response, timings=tuple(timings)) + response = replace(response, timings=tuple(timings), phrase_matches=tuple(phrase_matches)) logger.info( "ebook_search_complete vector_candidates=%s lexical_candidates=%s " - "fused_candidates=%s returned=%s rank_label=%s runtime_ms=%.1f", + "fused_candidates=%s phrase_matching=%s phrase_matches=%s returned=%s rank_label=%s runtime_ms=%.1f", len(retrieval.vector_results), len(retrieval.lexical_results), len(fused), + phrase_matching_enabled, + len(phrase_matches), len(response.results), response.rank_label, response.total_runtime_ms, @@ -130,6 +165,77 @@ def search_ebooks( return response +def skip_phrase_matches() -> list[HydratedPhraseMatch]: + """Return no protected phrase matches when phrase matching is disabled.""" + logger.info("ebook_protected_phrase_detection_skipped") + return [] + + +def query_phrase_matches(engine: Engine, query: str, config: EbookSearchConfig) -> list[HydratedPhraseMatch]: + """Detect protected phrases in a query without making search fail when phrase tables are unavailable.""" + try: + with Session(engine) as session: + return detect_protected_phrases_for_query(session, query, config) + except SQLAlchemyError as error: + logger.warning("ebook_protected_phrase_detection_unavailable error=%s", error) + return [] + + +def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchResult]: + """Return candidates unchanged when phrase matching is disabled.""" + logger.info("ebook_phrase_boost_skipped candidates=%s", len(candidates)) + return candidates + + +def apply_phrase_mention_boosts( + engine: Engine, + candidates: list[SearchResult], + phrase_matches: Sequence[HydratedPhraseMatch], + phrase_hit_boost: float, +) -> list[SearchResult]: + """Boost retrieved chunks that have indexed mentions for detected protected phrases.""" + phrase_ids = sorted({match.phrase_id for match in phrase_matches}) + if not candidates or not phrase_ids or phrase_hit_boost <= 0: + return candidates + + chunk_ids = [candidate.chunk_id for candidate in candidates] + try: + with Session(engine) as session: + hit_counts = phrase_hit_counts_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids) + except SQLAlchemyError as error: + logger.warning("ebook_phrase_boost_unavailable error=%s", error) + return candidates + + if not hit_counts: + return candidates + + boosted = [ + replace( + candidate, + score=candidate.score + (hit_counts.get(candidate.chunk_id, 0) * phrase_hit_boost), + fused_score=boosted_fused_score(candidate, hit_counts.get(candidate.chunk_id, 0), phrase_hit_boost), + phrase_hit_count=hit_counts.get(candidate.chunk_id, 0), + rank_source=phrase_rank_source(candidate.rank_source, hit_counts.get(candidate.chunk_id, 0)), + ) + for candidate in candidates + ] + return sorted(boosted, key=lambda candidate: candidate.score, reverse=True) + + +def boosted_fused_score(candidate: SearchResult, phrase_hit_count: int, phrase_hit_boost: float) -> float | None: + """Return a fused score adjusted by phrase hits when a fused score exists.""" + if candidate.fused_score is None: + return None + return candidate.fused_score + (phrase_hit_count * phrase_hit_boost) + + +def phrase_rank_source(rank_source: str, phrase_hit_count: int) -> str: + """Append phrase evidence to a rank-source label when a chunk was boosted.""" + if phrase_hit_count <= 0 or "phrases" in rank_source: + return rank_source + return f"{rank_source} + phrases" + + def parallel_retrieval( engine: Engine, query: str, @@ -223,6 +329,7 @@ def vector_candidates(engine: Engine, query: str, config: EbookSearchConfig) -> select( EbookChunk.id.label("chunk_id"), EbookChunk.text.label("text"), + EbookSource.id.label("source_id"), EbookSource.title.label("source_title"), EbookSource.author.label("source_author"), EbookChapter.title.label("chapter_title"), @@ -317,9 +424,11 @@ def reciprocal_rank_fusion( def search_result_from_row(row: Mapping[str, object]) -> SearchResult: """Convert a database row mapping into a search result.""" + source_id = row.get("source_id") return SearchResult( chunk_id=int(row["chunk_id"]), text=str(row["text"]), + source_id=int(source_id) if source_id is not None else None, source_title=str(row["source_title"]), source_author=optional_str(row["source_author"]), chapter_title=optional_str(row["chapter_title"]),