483 lines
17 KiB
Python
483 lines
17 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.chunk_records import CHUNK_RECORD_COLUMNS
|
|
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 PhraseMatch
|
|
|
|
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[PhraseMatch, ...] = ()
|
|
|
|
@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[PhraseMatch]
|
|
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)
|
|
phrase_boost_timing_name = "Phrase mention boost" if phrase_matching else "Phrase mention boost skipped"
|
|
fused, timing = await async_timed_result(
|
|
phrase_boost_timing_name,
|
|
apply_phrase_mention_boosts(
|
|
engine,
|
|
fused,
|
|
phrase_matches,
|
|
config.phrase_hit_boost,
|
|
phrase_matching=phrase_matching,
|
|
),
|
|
)
|
|
timings.append(timing)
|
|
rerank_enabled = config.rerank.enabled and rerank
|
|
rerank_timing_name = "Rerank" if rerank_enabled else "Rerank skipped"
|
|
response, timing = await async_timed_result(
|
|
rerank_timing_name,
|
|
apply_rerank(client, query, fused, config, rerank=rerank_enabled),
|
|
)
|
|
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
|
|
|
|
|
|
async def query_phrase_matches(
|
|
engine: AsyncEngine,
|
|
query: str,
|
|
config: EbookSearchConfig,
|
|
*,
|
|
phrase_matching: bool,
|
|
) -> list[PhraseMatch]:
|
|
"""Detect protected phrases in a query without making search fail when phrase tables are unavailable."""
|
|
if not phrase_matching:
|
|
logger.info("ebook_protected_phrase_detection_skipped")
|
|
return []
|
|
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 []
|
|
|
|
|
|
async def apply_phrase_mention_boosts(
|
|
engine: AsyncEngine,
|
|
candidates: list[SearchResult],
|
|
phrase_matches: Sequence[PhraseMatch],
|
|
phrase_hit_boost: float,
|
|
*,
|
|
phrase_matching: bool,
|
|
) -> list[SearchResult]:
|
|
"""Boost retrieved chunks that have indexed mentions for detected protected phrases when enabled."""
|
|
if not phrase_matching:
|
|
logger.info(f"ebook_phrase_boost_skipped candidates={len(candidates)}")
|
|
return candidates
|
|
|
|
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 and returns immediately when phrase matching is disabled.
|
|
"""
|
|
phrase_timing_name = "Protected phrase detection" if phrase_matching else "Protected phrase detection skipped"
|
|
(
|
|
(vector_results, vector_timing),
|
|
(lexical_results, lexical_timing),
|
|
(phrase_matches, phrase_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)),
|
|
async_timed_result(
|
|
phrase_timing_name,
|
|
query_phrase_matches(engine, query, config, phrase_matching=phrase_matching),
|
|
),
|
|
)
|
|
|
|
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),
|
|
),
|
|
)
|
|
|
|
|
|
async def apply_rerank(
|
|
client: httpx.AsyncClient,
|
|
query: str,
|
|
candidates: list[SearchResult],
|
|
config: EbookSearchConfig,
|
|
*,
|
|
rerank: bool,
|
|
) -> SearchResponse:
|
|
"""Rerank already-fused hybrid candidates when enabled for this request."""
|
|
if not rerank:
|
|
logger.info(f"ebook_rerank_skipped candidates={len(candidates)}")
|
|
return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
|
|
|
|
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(
|
|
*CHUNK_RECORD_COLUMNS,
|
|
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)
|