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"]),