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

feat(ebook): implement phrase matching functionality and UI enhancements

Browse Source
This commit is contained in:
Richie
2026-07-12 17:49:52 -04:00
parent ff0013c617
commit d03daeb66d
9 changed files with 151 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ebook_search/search.py
+113 -4
View File
@@ -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"]),