Richie
ed1ea4546a
Move protected phrase detection into the retrieval gather so it runs concurrently with vector and BM25 candidates instead of sequentially before them. Make the search API accept real bool form fields for rerank/phrase_matching, gate phrase matching on both the request and config kill switch, and reflow log f-strings for readability.
489 lines
18 KiB
Python
489 lines
18 KiB
Python
"""Hybrid search orchestration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import re
|
|
from collections import defaultdict
|
|
from dataclasses import dataclass, replace
|
|
from typing import TYPE_CHECKING
|
|
|
|
from pgvector.sqlalchemy import Vector
|
|
from sqlalchemy import literal, select
|
|
from sqlalchemy.exc import SQLAlchemyError
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from python.ebook_search.bm25_corpus import (
|
|
BM25CorpusUnavailableError,
|
|
load_bm25_corpus,
|
|
score_bm25_corpus,
|
|
)
|
|
from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table
|
|
from python.ebook_search.protected_phrases.matching import (
|
|
detect_protected_phrases_for_query,
|
|
phrase_hits_for_chunks,
|
|
)
|
|
from python.ebook_search.rerank import rerank_chunks
|
|
from python.ebook_search.timing import RuntimeStep, async_timed_result, timed_result
|
|
from python.orm.richie import (
|
|
EbookChapter,
|
|
EbookChunk,
|
|
EbookEmbeddingModel,
|
|
EbookSource,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Mapping, Sequence
|
|
|
|
import httpx
|
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
|
from python.ebook_search.config import EbookSearchConfig
|
|
from python.ebook_search.protected_phrases.models import HydratedPhraseMatch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SearchResult:
|
|
"""One source chunk returned by search."""
|
|
|
|
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
|
|
matched_phrases: tuple[str, ...] = ()
|
|
source_author: str | None = None
|
|
chapter_title: str | None = None
|
|
page_label: str | None = None
|
|
rank_source: str = "Hybrid"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SearchResponse:
|
|
"""Search output for the UI."""
|
|
|
|
query: str
|
|
results: list[SearchResult]
|
|
rank_label: str
|
|
timings: tuple[RuntimeStep, ...] = ()
|
|
phrase_matches: tuple[HydratedPhraseMatch, ...] = ()
|
|
|
|
@property
|
|
def total_runtime_ms(self) -> float:
|
|
"""Return total measured runtime for the response."""
|
|
return sum(step.duration_ms for step in self.timings if step.counts_toward_total)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RetrievalResponse:
|
|
"""Parallel retrieval output for vector, BM25, and protected phrase candidates."""
|
|
|
|
vector_results: list[SearchResult]
|
|
lexical_results: list[SearchResult]
|
|
phrase_matches: list[HydratedPhraseMatch]
|
|
timings: tuple[RuntimeStep, ...]
|
|
|
|
|
|
async def search_ebooks(
|
|
engine: AsyncEngine,
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
rerank: bool,
|
|
phrase_matching: bool,
|
|
) -> SearchResponse:
|
|
"""Run hybrid vector/BM25 search and optional reranking.
|
|
|
|
Phrase matching only runs when both the request asks for it and
|
|
``config.phrase_matching_enabled`` allows it.
|
|
"""
|
|
if not query.strip():
|
|
logger.info("ebook_search_empty_query")
|
|
return SearchResponse(query=query, results=[], rank_label="Hybrid")
|
|
|
|
phrase_matching = phrase_matching and config.phrase_matching_enabled
|
|
logger.info(f"ebook_search_start query_length={len(query)} {rerank=} {phrase_matching=}")
|
|
timings: list[RuntimeStep] = []
|
|
retrieval, timing = await async_timed_result(
|
|
"Hybrid retrieval",
|
|
parallel_retrieval(engine, client, query, config, phrase_matching=phrase_matching),
|
|
)
|
|
phrase_matches = retrieval.phrase_matches
|
|
timings.extend(retrieval.timings)
|
|
timings.append(timing)
|
|
fused, timing = timed_result(
|
|
"Reciprocal rank fusion",
|
|
reciprocal_rank_fusion,
|
|
retrieval.vector_results,
|
|
retrieval.lexical_results,
|
|
rank_constant=config.rrf_rank_constant,
|
|
)
|
|
timings.append(timing)
|
|
if phrase_matching:
|
|
fused, timing = await async_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 = await async_timed_result("Rerank", apply_rerank(client, query, fused, config))
|
|
else:
|
|
response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
|
|
timings.append(timing)
|
|
response = replace(response, timings=tuple(timings), phrase_matches=tuple(phrase_matches))
|
|
logger.info(
|
|
f"ebook_search_complete vector_candidates={len(retrieval.vector_results)} "
|
|
f"lexical_candidates={len(retrieval.lexical_results)} fused_candidates={len(fused)} {phrase_matching=} "
|
|
f"phrase_matches={len(phrase_matches)} returned={len(response.results)} {response.rank_label=} "
|
|
f"{response.total_runtime_ms=:.1f}"
|
|
)
|
|
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 []
|
|
|
|
|
|
async def query_phrase_matches(
|
|
engine: AsyncEngine,
|
|
query: str,
|
|
config: EbookSearchConfig,
|
|
) -> list[HydratedPhraseMatch]:
|
|
"""Detect protected phrases in a query without making search fail when phrase tables are unavailable."""
|
|
try:
|
|
async with AsyncSession(engine) as session:
|
|
return await detect_protected_phrases_for_query(session, query, config)
|
|
except SQLAlchemyError as error:
|
|
logger.warning(f"ebook_protected_phrase_detection_unavailable {error=}")
|
|
return []
|
|
|
|
|
|
def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchResult]:
|
|
"""Return candidates unchanged when phrase matching is disabled."""
|
|
logger.info(f"ebook_phrase_boost_skipped candidates={len(candidates)}")
|
|
return candidates
|
|
|
|
|
|
async def apply_phrase_mention_boosts(
|
|
engine: AsyncEngine,
|
|
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:
|
|
async with AsyncSession(engine) as session:
|
|
phrase_hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
|
|
except SQLAlchemyError as error:
|
|
logger.warning(f"ebook_phrase_boost_unavailable {error=}")
|
|
return candidates
|
|
|
|
if not phrase_hits:
|
|
return candidates
|
|
|
|
hit_counts = {
|
|
chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in phrase_hits.items()
|
|
}
|
|
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),
|
|
matched_phrases=tuple(hit.phrase_text for hit in phrase_hits.get(candidate.chunk_id, ())),
|
|
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"
|
|
|
|
|
|
async def parallel_retrieval(
|
|
engine: AsyncEngine,
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
phrase_matching: bool,
|
|
) -> RetrievalResponse:
|
|
"""Run vector, BM25, and protected phrase retrieval concurrently with separate database sessions.
|
|
|
|
BM25 scoring is pure CPU work over the cached corpus, so it runs in a worker thread
|
|
instead of on the event loop. Protected phrase detection only depends on the query, so
|
|
it joins the gather as a third task when phrase matching is enabled.
|
|
"""
|
|
phrase_task = (
|
|
asyncio.create_task(
|
|
async_timed_result("Protected phrase detection", query_phrase_matches(engine, query, config))
|
|
)
|
|
if phrase_matching
|
|
else None
|
|
)
|
|
(vector_results, vector_timing), (lexical_results, lexical_timing) = await asyncio.gather(
|
|
async_timed_result("Embedding + vector search", vector_candidates(engine, client, query, config)),
|
|
async_timed_result("BM25 search", asyncio.to_thread(bm25_candidates, query, config)),
|
|
)
|
|
if phrase_task is not None:
|
|
phrase_matches, phrase_timing = await phrase_task
|
|
else:
|
|
phrase_matches, phrase_timing = timed_result("Protected phrase detection skipped", skip_phrase_matches)
|
|
|
|
logger.info(
|
|
f"ebook_parallel_retrieval_complete vector_candidates={len(vector_results)} "
|
|
f"lexical_candidates={len(lexical_results)} phrase_matches={len(phrase_matches)}"
|
|
)
|
|
return RetrievalResponse(
|
|
vector_results=vector_results,
|
|
lexical_results=lexical_results,
|
|
phrase_matches=phrase_matches,
|
|
timings=(
|
|
replace(vector_timing, counts_toward_total=False),
|
|
replace(lexical_timing, counts_toward_total=False),
|
|
replace(phrase_timing, counts_toward_total=False),
|
|
),
|
|
)
|
|
|
|
|
|
def skip_rerank(
|
|
query: str,
|
|
candidates: list[SearchResult],
|
|
config: EbookSearchConfig,
|
|
) -> SearchResponse:
|
|
"""Return fused hybrid results without reranking."""
|
|
logger.info(f"ebook_rerank_skipped candidates={len(candidates)}")
|
|
return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
|
|
|
|
|
|
async def apply_rerank(
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
candidates: list[SearchResult],
|
|
config: EbookSearchConfig,
|
|
) -> SearchResponse:
|
|
"""Rerank already-fused hybrid candidates."""
|
|
reranked = await rerank_chunks(client, query, candidates[: config.rerank.candidates], config.rerank)
|
|
logger.info(
|
|
f"ebook_rerank_complete input_candidates={min(len(candidates), config.rerank.candidates)} "
|
|
f"returned={len(reranked)}"
|
|
)
|
|
return SearchResponse(
|
|
query=query,
|
|
results=[replace(result, rank_source="Hybrid + rerank") for result in reranked[: config.top_k]],
|
|
rank_label="Hybrid + rerank",
|
|
)
|
|
|
|
|
|
async def vector_candidates(
|
|
engine: AsyncEngine,
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
config: EbookSearchConfig,
|
|
) -> list[SearchResult]:
|
|
"""Return pgvector cosine candidates for a natural-language query."""
|
|
async with AsyncSession(engine) as session:
|
|
model = await session.scalar(
|
|
select(EbookEmbeddingModel).where(EbookEmbeddingModel.name == config.embedding_model)
|
|
)
|
|
if model is None:
|
|
msg = f"Embedding model is not registered: {config.embedding_model}"
|
|
raise ValueError(msg)
|
|
|
|
expected_dimension = MODEL_DIMENSIONS[config.embedding_model]
|
|
if model.dimension != expected_dimension:
|
|
msg = f"Model row dimension {model.dimension} does not match configured dimension {expected_dimension}"
|
|
raise ValueError(msg)
|
|
|
|
embedding = await embed_query(client, query, config)
|
|
limit = max(config.rerank.candidates, config.top_k) * config.vector_candidate_multiplier
|
|
embedding_table = get_embedding_table(model.dimension)
|
|
|
|
embedding_param = literal(embedding, type_=Vector(model.dimension))
|
|
distance = embedding_table.embedding.op("<=>")(embedding_param)
|
|
score = (literal(1.0) - distance).label("score")
|
|
statement = (
|
|
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"),
|
|
EbookChunk.page_label.label("page_label"),
|
|
score,
|
|
)
|
|
.select_from(embedding_table)
|
|
.join(EbookChunk, EbookChunk.id == embedding_table.chunk_id)
|
|
.join(EbookSource, EbookSource.id == EbookChunk.source_id)
|
|
.outerjoin(EbookChapter, EbookChapter.id == EbookChunk.chapter_id)
|
|
.where(embedding_table.model_id == model.id)
|
|
.order_by(distance)
|
|
.limit(limit)
|
|
)
|
|
rows = (await session.execute(statement)).mappings()
|
|
results = [search_result_from_row(row) for row in rows]
|
|
logger.info(
|
|
f"ebook_vector_search_complete {config.embedding_model=} {model.dimension=} candidates={len(results)}"
|
|
)
|
|
return results
|
|
|
|
|
|
def bm25_candidates(query: str, config: EbookSearchConfig) -> list[SearchResult]:
|
|
"""Return BM25-ranked lexical candidates using the persisted corpus."""
|
|
try:
|
|
corpus = load_bm25_corpus(config)
|
|
except BM25CorpusUnavailableError as error:
|
|
logger.warning(f"ebook_bm25_index_unavailable_skipping {error=}")
|
|
return []
|
|
|
|
if not corpus.records:
|
|
logger.info("ebook_bm25_search_complete corpus=0 candidates=0")
|
|
return []
|
|
|
|
bm25_query = retrieval_query_from_text(query)
|
|
scored_records = score_bm25_corpus(bm25_query, corpus, limit=config.bm25_candidate_limit)
|
|
results = [
|
|
replace(search_result_from_row(record), score=score, vector_score=None, bm25_score=score)
|
|
for record, score in scored_records
|
|
]
|
|
|
|
max_score = results[0].bm25_score if results else 0.0
|
|
logger.info(f"ebook_bm25_search_complete corpus={len(corpus.records)} candidates={len(results)} {max_score=:.6f}")
|
|
return results
|
|
|
|
|
|
def reciprocal_rank_fusion(
|
|
vector_results: list[SearchResult],
|
|
lexical_results: list[SearchResult],
|
|
rank_constant: int,
|
|
) -> list[SearchResult]:
|
|
"""Fuse vector and lexical rankings with Reciprocal Rank Fusion."""
|
|
by_chunk: dict[int, SearchResult] = {}
|
|
scores: defaultdict[int, float] = defaultdict(float)
|
|
vector_scores: dict[int, float] = {}
|
|
bm25_scores: dict[int, float] = {}
|
|
|
|
for rank, result in enumerate(vector_results, start=1):
|
|
by_chunk.setdefault(result.chunk_id, result)
|
|
vector_scores[result.chunk_id] = result.vector_score if result.vector_score is not None else result.score
|
|
scores[result.chunk_id] += 1 / (rank_constant + rank)
|
|
|
|
for rank, result in enumerate(lexical_results, start=1):
|
|
by_chunk.setdefault(result.chunk_id, result)
|
|
bm25_scores[result.chunk_id] = result.bm25_score if result.bm25_score is not None else result.score
|
|
scores[result.chunk_id] += 1 / (rank_constant + rank)
|
|
|
|
return sorted(
|
|
(
|
|
replace(
|
|
result,
|
|
score=scores[result.chunk_id],
|
|
vector_score=vector_scores.get(result.chunk_id),
|
|
bm25_score=bm25_scores.get(result.chunk_id),
|
|
fused_score=scores[result.chunk_id],
|
|
rank_source="Hybrid",
|
|
)
|
|
for result in by_chunk.values()
|
|
),
|
|
key=lambda result: result.score,
|
|
reverse=True,
|
|
)
|
|
|
|
|
|
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"]),
|
|
page_label=optional_str(row["page_label"]),
|
|
score=float(row["score"]) if "score" in row else 0.0,
|
|
vector_score=float(row["score"]) if "score" in row else None,
|
|
)
|
|
|
|
|
|
def optional_str(value: object) -> str | None:
|
|
"""Convert nullable database values to optional strings."""
|
|
if value is None:
|
|
return None
|
|
return str(value)
|
|
|
|
|
|
TOKEN_RE = re.compile(r"[A-Za-z0-9_]+")
|
|
|
|
|
|
def tokens(text_value: str) -> list[str]:
|
|
"""Extract tokens from a text value.
|
|
|
|
This is a simple approximation of the tokenization used by PostgreSQL's full-text search,
|
|
which is sufficient for BM25 candidate retrieval. It lowercases tokens and includes alphanumeric characters and
|
|
underscores.
|
|
"""
|
|
return [match.group(0).lower() for match in TOKEN_RE.finditer(text_value)]
|
|
|
|
|
|
QUERY_STOP_WORDS = {
|
|
"a",
|
|
"an",
|
|
"and",
|
|
"are",
|
|
"as",
|
|
"at",
|
|
"does",
|
|
"for",
|
|
"in",
|
|
"is",
|
|
"of",
|
|
"the",
|
|
"to",
|
|
"what",
|
|
"when",
|
|
"where",
|
|
"which",
|
|
"who",
|
|
"why",
|
|
}
|
|
|
|
|
|
def retrieval_query_from_text(query: str) -> str:
|
|
"""Remove generic question words while preserving entity and series terms."""
|
|
keywords = [token for token in tokens(query) if token not in QUERY_STOP_WORDS]
|
|
if not keywords:
|
|
return query
|
|
return " ".join(keywords)
|