refactor(ebook-search): simplify search and phrase matching
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 28s
test ebook search / test-ebook-search (pull_request) Failing after 35s
build_systems / build-bob (pull_request) Successful in 51s
build_systems / build-brain (pull_request) Successful in 50s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m3s
build_systems / build-jeeves (pull_request) Successful in 2m20s
treefmt / nix fmt (pull_request) Successful in 5s
pytest / pytest (pull_request) Successful in 28s
test ebook search / test-ebook-search (pull_request) Failing after 35s
build_systems / build-bob (pull_request) Successful in 51s
build_systems / build-brain (pull_request) Successful in 50s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m3s
build_systems / build-jeeves (pull_request) Successful in 2m20s
This commit is contained in:
@@ -13,7 +13,7 @@ from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolv
|
||||
AppEngine,
|
||||
AppHttpClient,
|
||||
)
|
||||
from python.ebook_search.api.web import templates
|
||||
from python.ebook_search.api.web import error_response, templates
|
||||
from python.ebook_search.embeddings import embed_missing_chunks, embedding_model_stats
|
||||
from python.ebook_search.ingest import ingest_configured_paths
|
||||
from python.ebook_search.protected_phrases.generate_ngrams import generate_candidate_phrases_for_books
|
||||
@@ -50,7 +50,7 @@ async def scan_library(request: Request, config: AppConfig, session: AsyncDbSess
|
||||
await session.commit()
|
||||
except Exception as error:
|
||||
logger.exception("ebook_admin_scan_failed")
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
return error_response(request, error)
|
||||
|
||||
logger.info(f"ebook_admin_scan_complete {count=}")
|
||||
if count > 0:
|
||||
@@ -65,7 +65,7 @@ async def generate_all_phrases(request: Request, config: AppConfig, engine: AppE
|
||||
result = await generate_candidate_phrases_for_books(engine, config)
|
||||
except Exception as error:
|
||||
logger.exception("ebook_admin_generate_phrases_failed")
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
return error_response(request, error)
|
||||
|
||||
logger.info(
|
||||
f"ebook_admin_generate_phrases_complete {result.books_seen=} {result.books_built=} {result.candidate_phrases=}"
|
||||
@@ -128,7 +128,7 @@ async def run_phrase_judgment(
|
||||
result = await judge_candidate_phrases_for_books(engine, config, source_ids=source_ids)
|
||||
except Exception as error:
|
||||
logger.exception("ebook_admin_judge_phrases_failed")
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
return error_response(request, error)
|
||||
|
||||
logger.info(
|
||||
f"ebook_admin_judge_phrases_complete {result.books_seen=} {result.books_judged=} {result.books_failed=} "
|
||||
@@ -161,7 +161,7 @@ async def embed_missing(
|
||||
await session.commit()
|
||||
except Exception as error:
|
||||
logger.exception("ebook_admin_embed_missing_failed")
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
return error_response(request, error)
|
||||
|
||||
logger.info(f"ebook_admin_embed_missing_complete {count=}")
|
||||
return templates.TemplateResponse(
|
||||
@@ -192,12 +192,7 @@ async def embed_all(
|
||||
logger.info(f"ebook_admin_embed_all_batch_complete {batches=} {count=} {total=}")
|
||||
except Exception as error:
|
||||
logger.exception(f"ebook_admin_embed_all_failed {batches=} {total=}")
|
||||
return templates.TemplateResponse(
|
||||
request,
|
||||
"partials/error.html",
|
||||
{"message": f"Embed all failed after {total} chunks in {batches} batches: {error}"},
|
||||
status_code=500,
|
||||
)
|
||||
return error_response(request, f"Embed all failed after {total} chunks in {batches} batches: {error}")
|
||||
|
||||
logger.info(f"ebook_admin_embed_all_complete {batches=} {total=}")
|
||||
return templates.TemplateResponse(
|
||||
|
||||
@@ -15,6 +15,7 @@ from python.ebook_search.api.dependencies import (
|
||||
from python.ebook_search.api.judge_tasks import is_judging_book, pop_book_judgment_outcome, start_book_phrase_judgment
|
||||
from python.ebook_search.api.web import templates
|
||||
from python.ebook_search.protected_phrases.generate_ngrams import recalculate_candidate_phrases_for_book
|
||||
from python.ebook_search.protected_phrases.store import count_protected_phrases
|
||||
from python.fastapi_tools import AsyncDbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime
|
||||
from python.orm.richie import EbookCandidatePhrase, EbookChapter, EbookChunk, EbookProtectedPhrase, EbookSource
|
||||
|
||||
@@ -71,14 +72,6 @@ async def get_judged_candidate_count(session: AsyncSession, book_id: int) -> int
|
||||
)
|
||||
|
||||
|
||||
async def get_protected_count(session: AsyncSession, book_id: int) -> int:
|
||||
"""Return the number of protected phrases for one book."""
|
||||
return (
|
||||
await session.scalar(select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id))
|
||||
or 0
|
||||
)
|
||||
|
||||
|
||||
async def get_candidates(session: AsyncSession, book_id: int) -> list[EbookCandidatePhrase]:
|
||||
"""Return the indexed candidates for one book."""
|
||||
return list(
|
||||
@@ -122,7 +115,7 @@ async def book_detail(source_id: int, request: Request, session: AsyncDbSession)
|
||||
chunk_count = await get_chunk_count(session, source.id)
|
||||
candidate_count = await get_candidate_count(session, source.id)
|
||||
judged_candidate_count = await get_judged_candidate_count(session, source.id)
|
||||
protected_count = await get_protected_count(session, source.id)
|
||||
protected_count = await count_protected_phrases(session, source.id)
|
||||
candidates = await get_candidates(session, source.id)
|
||||
protected_phrases = await get_protected_phrases(session, source.id)
|
||||
else:
|
||||
|
||||
@@ -16,7 +16,7 @@ from python.ebook_search.api.dependencies import ( # noqa: TC001 FastAPI resolv
|
||||
AppEngine,
|
||||
AppHttpClient,
|
||||
)
|
||||
from python.ebook_search.api.web import templates
|
||||
from python.ebook_search.api.web import error_response, templates
|
||||
from python.ebook_search.guardrails import (
|
||||
CitationReport,
|
||||
is_confident,
|
||||
@@ -95,7 +95,7 @@ async def search(
|
||||
)
|
||||
except Exception as error:
|
||||
logger.exception("ebook_search_request_failed")
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500)
|
||||
return error_response(request, error)
|
||||
|
||||
answer_start = perf_counter()
|
||||
answer, low_confidence, citation_report = await build_answer(client, query, response, config)
|
||||
|
||||
@@ -3,9 +3,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from fastapi.templating import Jinja2Templates
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastapi import Request
|
||||
from fastapi.responses import HTMLResponse
|
||||
|
||||
PACKAGE_DIR = Path(__file__).resolve().parent
|
||||
TEMPLATE_DIR = PACKAGE_DIR / "templates"
|
||||
STATIC_DIR = PACKAGE_DIR / "static"
|
||||
@@ -21,3 +26,8 @@ def static_version(filename: str) -> int:
|
||||
|
||||
templates = Jinja2Templates(directory=TEMPLATE_DIR)
|
||||
templates.env.globals["static_version"] = static_version
|
||||
|
||||
|
||||
def error_response(request: Request, message: object) -> HTMLResponse:
|
||||
"""Render the shared error partial for a failed UI request."""
|
||||
return templates.TemplateResponse(request, "partials/error.html", {"message": str(message)}, status_code=500)
|
||||
|
||||
@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING
|
||||
import bm25s
|
||||
from sqlalchemy import func, select, union_all
|
||||
|
||||
from python.ebook_search.chunk_records import CHUNK_RECORD_COLUMNS
|
||||
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -169,13 +170,7 @@ async def fetch_bm25_corpus_records(session: AsyncSession) -> tuple[list[dict[st
|
||||
"""
|
||||
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"),
|
||||
*CHUNK_RECORD_COLUMNS,
|
||||
EbookChunk.search_text.label("bm25_text"),
|
||||
)
|
||||
.select_from(EbookChunk)
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Shared database columns used to build search-result records."""
|
||||
|
||||
from python.orm.richie import EbookChapter, EbookChunk, EbookSource
|
||||
|
||||
CHUNK_RECORD_COLUMNS = (
|
||||
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"),
|
||||
)
|
||||
@@ -2,7 +2,6 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from os import getenv
|
||||
from typing import Annotated, Self
|
||||
|
||||
from pydantic import AliasChoices, Field, field_validator, model_validator
|
||||
@@ -32,11 +31,6 @@ def normalize_embedding_alias(model: str) -> str:
|
||||
return standard_model
|
||||
|
||||
|
||||
def normalize_embedding_model(default: str = "qwen3-embedding-0.6b") -> str:
|
||||
"""Normalize the configured embedding alias to its provider model name."""
|
||||
return normalize_embedding_alias(getenv("EBOOK_SEARCH_EMBEDDING_MODEL", default))
|
||||
|
||||
|
||||
class RerankConfig(BaseSettings):
|
||||
"""vLLM reranker settings."""
|
||||
|
||||
|
||||
@@ -64,17 +64,13 @@ async def check_embedding_endpoint(
|
||||
timeout_seconds: float = 5.0,
|
||||
) -> bool:
|
||||
"""Return whether the configured embedding endpoint answers a model listing."""
|
||||
try:
|
||||
response = await client.get(
|
||||
f"{config.embedding_base_url.rstrip('/')}/models",
|
||||
headers=auth_headers(config.embedding_api_key),
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as error:
|
||||
logger.warning(f"ebook_embedding_endpoint_unreachable {config.embedding_base_url=} {error=}")
|
||||
return False
|
||||
return True
|
||||
return await _check_endpoint(
|
||||
client,
|
||||
base_url=config.embedding_base_url,
|
||||
api_key=config.embedding_api_key,
|
||||
timeout_seconds=timeout_seconds,
|
||||
unavailable_log=f"ebook_embedding_endpoint_unreachable {config.embedding_base_url=}",
|
||||
)
|
||||
|
||||
|
||||
async def check_chat_endpoint(
|
||||
@@ -84,15 +80,33 @@ async def check_chat_endpoint(
|
||||
timeout_seconds: float = 5.0,
|
||||
) -> bool:
|
||||
"""Return whether the configured chat (answering) endpoint answers a model listing."""
|
||||
return await _check_endpoint(
|
||||
client,
|
||||
base_url=config.vllm_base_url,
|
||||
api_key=config.vllm_api_key,
|
||||
timeout_seconds=timeout_seconds,
|
||||
unavailable_log=f"ebook_chat_endpoint_unreachable {config.vllm_base_url=}",
|
||||
)
|
||||
|
||||
|
||||
async def _check_endpoint(
|
||||
client: httpx.AsyncClient,
|
||||
*,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
timeout_seconds: float,
|
||||
unavailable_log: str,
|
||||
) -> bool:
|
||||
"""Return whether an OpenAI-compatible endpoint answers a model listing."""
|
||||
try:
|
||||
response = await client.get(
|
||||
f"{config.vllm_base_url.rstrip('/')}/models",
|
||||
headers=auth_headers(config.vllm_api_key),
|
||||
f"{base_url.rstrip('/')}/models",
|
||||
headers=auth_headers(api_key),
|
||||
timeout=timeout_seconds,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPError as error:
|
||||
logger.warning(f"ebook_chat_endpoint_unreachable {config.vllm_base_url=} {error=}")
|
||||
logger.warning(f"{unavailable_log} {error=}")
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
@@ -64,40 +64,30 @@ def normalize_candidate_phrase(
|
||||
phrase_text: str,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
min_tokens: int | None = None,
|
||||
max_tokens: int | None = None,
|
||||
strip_leading_article: bool = False,
|
||||
) -> tuple[str, str, int] | None:
|
||||
"""Normalize a candidate phrase and validate token bounds.
|
||||
|
||||
Args:
|
||||
phrase_text (str): Raw phrase text to normalize.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
min_tokens (int | None): Minimum token count override; defaults to ``config.phrase_min_tokens``.
|
||||
max_tokens (int | None): Maximum token count override; defaults to ``config.phrase_max_tokens``.
|
||||
strip_leading_article (bool): Whether to drop a single leading English article.
|
||||
|
||||
Returns:
|
||||
tuple[str, str, int] | None: Display text, normalized phrase, and token count, or ``None``
|
||||
when the phrase falls outside the token bounds or is ignored.
|
||||
"""
|
||||
normalized_tokens = tokenize_with_offsets(phrase_text)
|
||||
start = 0
|
||||
if strip_leading_article and normalized_tokens and normalized_tokens[0].text in {"the", "a", "an"}:
|
||||
start = 1
|
||||
|
||||
selected_tokens = normalized_tokens[start:]
|
||||
min_count = config.phrase_min_tokens if min_tokens is None else min_tokens
|
||||
max_count = config.phrase_max_tokens if max_tokens is None else max_tokens
|
||||
if len(selected_tokens) < min_count or len(selected_tokens) > max_count:
|
||||
if len(normalized_tokens) < config.phrase_min_tokens or len(normalized_tokens) > max_count:
|
||||
return None
|
||||
|
||||
phrase_norm = " ".join(token.text for token in selected_tokens)
|
||||
phrase_norm = " ".join(token.text for token in normalized_tokens)
|
||||
if phrase_norm in get_ignored_phrases():
|
||||
return None
|
||||
|
||||
display_text = phrase_text[selected_tokens[0].start_char : selected_tokens[-1].end_char].strip()
|
||||
return display_text or phrase_norm, phrase_norm, len(selected_tokens)
|
||||
display_text = phrase_text[normalized_tokens[0].start_char : normalized_tokens[-1].end_char].strip()
|
||||
return display_text or phrase_norm, phrase_norm, len(normalized_tokens)
|
||||
|
||||
|
||||
def count_raw_ngrams(tokens: Sequence[str], config: EbookSearchConfig) -> Counter[str]:
|
||||
|
||||
@@ -6,12 +6,10 @@ import logging
|
||||
from collections import defaultdict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import and_, delete, func, or_, select
|
||||
from sqlalchemy import and_, delete, or_, select, union
|
||||
|
||||
from python.ebook_search.protected_phrases.config import get_ignored_phrases
|
||||
from python.ebook_search.protected_phrases.models import (
|
||||
ChunkPhraseHit,
|
||||
HydratedPhraseMatch,
|
||||
PhraseLookup,
|
||||
PhraseMatch,
|
||||
)
|
||||
@@ -29,11 +27,107 @@ if TYPE_CHECKING:
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
from python.ebook_search.protected_phrases.text_normalization import NormalizedToken
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def detect_protected_phrases_for_query(
|
||||
session: AsyncSession,
|
||||
query_text: str,
|
||||
config: EbookSearchConfig,
|
||||
) -> list[PhraseMatch]:
|
||||
"""Find query phrases with indexed exact matches on canonical and alias norms.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
query_text (str): User query text to detect phrases in.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: Metadata-backed, overlap-resolved phrase matches for the query.
|
||||
"""
|
||||
tokens_ = tokenize_with_offsets(query_text)
|
||||
windows_by_norm: defaultdict[str, list[tuple[int, int]]] = defaultdict(list)
|
||||
token_texts = [token.text for token in tokens_]
|
||||
max_tokens = max(config.phrase_max_tokens, config.phrase_max_entity_tokens)
|
||||
for phrase_norm, start, end in generate_query_ngrams(
|
||||
token_texts,
|
||||
min_n=config.phrase_min_tokens,
|
||||
max_n=max_tokens,
|
||||
):
|
||||
windows_by_norm[phrase_norm].append((start, end))
|
||||
if not windows_by_norm:
|
||||
return []
|
||||
|
||||
query_norms = tuple(windows_by_norm)
|
||||
matched_norms = union(
|
||||
select(
|
||||
EbookProtectedPhrase.id.label("phrase_id"),
|
||||
EbookProtectedPhrase.phrase_norm.label("matched_norm"),
|
||||
).where(EbookProtectedPhrase.phrase_norm.in_(query_norms)),
|
||||
select(
|
||||
EbookPhraseAlias.phrase_id.label("phrase_id"),
|
||||
EbookPhraseAlias.alias_norm.label("matched_norm"),
|
||||
).where(EbookPhraseAlias.alias_norm.in_(query_norms)),
|
||||
).subquery()
|
||||
statement = select(EbookProtectedPhrase, matched_norms.c.matched_norm).join(
|
||||
matched_norms,
|
||||
matched_norms.c.phrase_id == EbookProtectedPhrase.id,
|
||||
)
|
||||
|
||||
matches: list[PhraseMatch] = []
|
||||
for phrase, matched_norm in await session.execute(statement):
|
||||
for start, end in windows_by_norm[matched_norm]:
|
||||
matches.append(
|
||||
PhraseMatch(
|
||||
phrase_id=phrase.id,
|
||||
matched_norm=matched_norm,
|
||||
phrase_text=phrase.phrase_text,
|
||||
phrase_norm=phrase.phrase_norm,
|
||||
canonical_id=phrase.canonical_id,
|
||||
phrase_type=phrase.phrase_type,
|
||||
token_count=end - start,
|
||||
confidence=phrase.confidence,
|
||||
importance=phrase.importance,
|
||||
allow_nested=phrase.allow_nested,
|
||||
suppress_children=phrase.suppress_children,
|
||||
start_token=start,
|
||||
end_token=end,
|
||||
start_char=tokens_[start].start_char,
|
||||
end_char=tokens_[end - 1].end_char,
|
||||
book_id=phrase.book_id,
|
||||
series_id=phrase.series_id,
|
||||
)
|
||||
)
|
||||
return resolve_overlaps(matches)
|
||||
|
||||
|
||||
async def index_chunk_phrase_mentions_for_book(
|
||||
session: AsyncSession,
|
||||
book_id: int,
|
||||
config: EbookSearchConfig,
|
||||
) -> int:
|
||||
"""Rebuild chunk phrase mentions for all chunks in one book.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
book_id (int): Book whose chunk mentions are rebuilt.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
int: Total number of chunk phrase mentions indexed for the book.
|
||||
"""
|
||||
lookup = await load_phrase_lookup(session, config, book_id=book_id)
|
||||
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
|
||||
chunks = await session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id))
|
||||
count = 0
|
||||
for chunk in chunks:
|
||||
count += index_chunk_phrase_mentions(session, chunk, lookup=lookup)
|
||||
await session.flush()
|
||||
logger.info(f"ebook_chunk_phrase_mentions_indexed {book_id=} {count=}")
|
||||
return count
|
||||
|
||||
|
||||
async def load_phrase_lookup(
|
||||
session: AsyncSession,
|
||||
config: EbookSearchConfig,
|
||||
@@ -52,40 +146,29 @@ async def load_phrase_lookup(
|
||||
Returns:
|
||||
PhraseLookup: Normalized phrase and alias maps with the token-window bounds to test.
|
||||
"""
|
||||
norm_to_ids: defaultdict[str, list[int]] = defaultdict(list)
|
||||
alias_to_ids: defaultdict[str, list[int]] = defaultdict(list)
|
||||
phrase_ids_by_norm: defaultdict[str, set[int]] = defaultdict(set)
|
||||
phrases_by_id: dict[int, EbookProtectedPhrase] = {}
|
||||
max_tokens = config.phrase_max_tokens
|
||||
|
||||
phrase_statement = select(
|
||||
EbookProtectedPhrase.id,
|
||||
EbookProtectedPhrase.phrase_norm,
|
||||
EbookProtectedPhrase.token_count,
|
||||
)
|
||||
statement = select(
|
||||
EbookProtectedPhrase,
|
||||
EbookPhraseAlias.alias_norm,
|
||||
).outerjoin(EbookPhraseAlias, EbookPhraseAlias.phrase_id == EbookProtectedPhrase.id)
|
||||
scope_filter = protected_phrase_scope_filter(book_id=book_id, series_id=series_id)
|
||||
if scope_filter is not None:
|
||||
phrase_statement = phrase_statement.where(scope_filter)
|
||||
statement = statement.where(scope_filter)
|
||||
|
||||
for row in await session.execute(phrase_statement):
|
||||
phrase_id = int(row.id)
|
||||
phrase_norm = str(row.phrase_norm)
|
||||
norm_to_ids[phrase_norm].append(phrase_id)
|
||||
max_tokens = max(max_tokens, int(row.token_count))
|
||||
|
||||
alias_statement = select(
|
||||
EbookPhraseAlias.alias_norm,
|
||||
EbookPhraseAlias.phrase_id,
|
||||
).join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookPhraseAlias.phrase_id)
|
||||
if scope_filter is not None:
|
||||
alias_statement = alias_statement.where(scope_filter)
|
||||
|
||||
for row in await session.execute(alias_statement):
|
||||
alias_norm = str(row.alias_norm)
|
||||
alias_to_ids[alias_norm].append(int(row.phrase_id))
|
||||
max_tokens = max(max_tokens, len(alias_norm.split()))
|
||||
for phrase, alias_norm in await session.execute(statement):
|
||||
phrases_by_id[phrase.id] = phrase
|
||||
phrase_ids_by_norm[phrase.phrase_norm].add(phrase.id)
|
||||
max_tokens = max(max_tokens, phrase.token_count)
|
||||
if alias_norm is not None:
|
||||
phrase_ids_by_norm[alias_norm].add(phrase.id)
|
||||
max_tokens = max(max_tokens, len(alias_norm.split()))
|
||||
|
||||
return PhraseLookup(
|
||||
norm_to_phrase_ids={key: tuple(values) for key, values in norm_to_ids.items()},
|
||||
alias_to_phrase_ids={key: tuple(values) for key, values in alias_to_ids.items()},
|
||||
phrase_ids_by_norm={key: tuple(sorted(values)) for key, values in phrase_ids_by_norm.items()},
|
||||
phrases_by_id=phrases_by_id,
|
||||
min_tokens=config.phrase_min_tokens,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
@@ -111,6 +194,144 @@ def protected_phrase_scope_filter(*, book_id: int | None, series_id: int | None)
|
||||
return and_(*conditions)
|
||||
|
||||
|
||||
def is_inside(child: PhraseMatch, parent: PhraseMatch) -> bool:
|
||||
"""Return whether one token span is strictly inside another.
|
||||
|
||||
Args:
|
||||
child (PhraseMatch): Candidate nested match.
|
||||
parent (PhraseMatch): Candidate enclosing match.
|
||||
|
||||
Returns:
|
||||
bool: True when ``child`` lies within ``parent`` and is not the same span.
|
||||
"""
|
||||
return (
|
||||
child.start_token >= parent.start_token
|
||||
and child.end_token <= parent.end_token
|
||||
and (child.start_token, child.end_token, child.phrase_id)
|
||||
!= (parent.start_token, parent.end_token, parent.phrase_id)
|
||||
)
|
||||
|
||||
|
||||
def index_chunk_phrase_mentions(session: AsyncSession, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
|
||||
"""Store protected phrase mentions for one chunk.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
chunk (EbookChunk): Chunk whose text is scanned for phrase mentions.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
int: Number of phrase mentions stored for the chunk.
|
||||
"""
|
||||
tokens_ = tokenize_with_offsets(chunk.text)
|
||||
token_texts = [token.text for token in tokens_]
|
||||
raw_matches: list[PhraseMatch] = []
|
||||
phrase_windows = generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens)
|
||||
for matched_norm, start, end in phrase_windows:
|
||||
for phrase_id in lookup.phrase_ids_by_norm.get(matched_norm, ()):
|
||||
phrase = lookup.phrases_by_id[phrase_id]
|
||||
raw_matches.append(
|
||||
PhraseMatch(
|
||||
phrase_id=phrase_id,
|
||||
matched_norm=matched_norm,
|
||||
phrase_text=phrase.phrase_text,
|
||||
phrase_norm=phrase.phrase_norm,
|
||||
canonical_id=phrase.canonical_id,
|
||||
phrase_type=phrase.phrase_type,
|
||||
confidence=phrase.confidence,
|
||||
importance=phrase.importance,
|
||||
allow_nested=phrase.allow_nested,
|
||||
suppress_children=phrase.suppress_children,
|
||||
start_token=start,
|
||||
end_token=end,
|
||||
token_count=end - start,
|
||||
start_char=tokens_[start].start_char,
|
||||
end_char=tokens_[end - 1].end_char,
|
||||
book_id=phrase.book_id,
|
||||
series_id=phrase.series_id,
|
||||
)
|
||||
)
|
||||
|
||||
matches = resolve_overlaps(raw_matches)
|
||||
for match in matches:
|
||||
session.add(
|
||||
EbookChunkPhraseMention(
|
||||
chunk_id=chunk.id,
|
||||
phrase_id=match.phrase_id,
|
||||
book_id=match.book_id if match.book_id is not None else chunk.source_id,
|
||||
series_id=match.series_id,
|
||||
start_char=match.start_char if match.start_char is not None else 0,
|
||||
end_char=match.end_char,
|
||||
)
|
||||
)
|
||||
return len(matches)
|
||||
|
||||
|
||||
def resolve_overlaps(matches: Sequence[PhraseMatch]) -> list[PhraseMatch]:
|
||||
"""Resolve overlapping phrase matches without relying only on longest match.
|
||||
|
||||
Args:
|
||||
matches (Sequence[PhraseMatch]): Metadata-backed matches that may overlap.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: The kept, non-suppressed matches.
|
||||
"""
|
||||
sorted_matches = sorted(
|
||||
matches,
|
||||
key=lambda match: (match.start_token, -match.token_count, -match.importance, -match.confidence),
|
||||
)
|
||||
kept: list[PhraseMatch] = []
|
||||
for candidate in sorted_matches:
|
||||
if any(should_suppress(candidate, existing) for existing in kept):
|
||||
continue
|
||||
kept.append(candidate)
|
||||
return kept
|
||||
|
||||
|
||||
def should_suppress(candidate: PhraseMatch, kept: PhraseMatch) -> bool:
|
||||
"""Return whether an already-kept match should suppress a candidate.
|
||||
|
||||
Args:
|
||||
candidate (PhraseMatch): Match being considered for keeping.
|
||||
kept (PhraseMatch): Match already kept that may suppress the candidate.
|
||||
|
||||
Returns:
|
||||
bool: True when the candidate should be dropped in favor of the kept match.
|
||||
"""
|
||||
if not overlaps(candidate, kept):
|
||||
return False
|
||||
if candidate.canonical_id == kept.canonical_id:
|
||||
return rank_match(kept) >= rank_match(candidate)
|
||||
if is_inside(candidate, kept) and kept.suppress_children and not candidate.allow_nested:
|
||||
return True
|
||||
return not candidate.allow_nested and rank_match(kept) > rank_match(candidate)
|
||||
|
||||
|
||||
def overlaps(first: PhraseMatch, second: PhraseMatch) -> bool:
|
||||
"""Return whether two token spans overlap.
|
||||
|
||||
Args:
|
||||
first (PhraseMatch): First match to compare.
|
||||
second (PhraseMatch): Second match to compare.
|
||||
|
||||
Returns:
|
||||
bool: True when the two token spans share at least one token position.
|
||||
"""
|
||||
return not (first.end_token <= second.start_token or first.start_token >= second.end_token)
|
||||
|
||||
|
||||
def rank_match(match: PhraseMatch) -> tuple[float, float, int]:
|
||||
"""Rank phrase matches by importance, confidence, then token count.
|
||||
|
||||
Args:
|
||||
match (PhraseMatch): Match to build a sort key for.
|
||||
|
||||
Returns:
|
||||
tuple[float, float, int]: A comparable key of importance, confidence, and token count.
|
||||
"""
|
||||
return (match.importance, match.confidence, match.token_count)
|
||||
|
||||
|
||||
def generate_query_ngrams(
|
||||
tokens_: Sequence[str],
|
||||
min_n: int,
|
||||
@@ -134,358 +355,3 @@ def generate_query_ngrams(
|
||||
if phrase_norm in get_ignored_phrases():
|
||||
continue
|
||||
yield phrase_norm, start, end
|
||||
|
||||
|
||||
def detect_phrase_candidates(query_text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
|
||||
"""Detect protected phrase windows in a user query using RAM hash lookups.
|
||||
|
||||
Args:
|
||||
query_text (str): User query text to scan.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: Unhydrated phrase matches found in the query.
|
||||
"""
|
||||
return detect_phrase_candidates_from_tokens(tokenize_with_offsets(query_text), lookup)
|
||||
|
||||
|
||||
def detect_phrase_candidates_in_text(text: str, lookup: PhraseLookup) -> list[PhraseMatch]:
|
||||
"""Detect protected phrase windows in arbitrary text with character offsets.
|
||||
|
||||
Args:
|
||||
text (str): Arbitrary text, such as a chunk, to scan.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: Unhydrated phrase matches found in the text.
|
||||
"""
|
||||
return detect_phrase_candidates_from_tokens(tokenize_with_offsets(text), lookup)
|
||||
|
||||
|
||||
def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], lookup: PhraseLookup) -> list[PhraseMatch]:
|
||||
"""Detect protected phrase windows from already-normalized tokens.
|
||||
|
||||
Args:
|
||||
tokens_ (Sequence[NormalizedToken]): Normalized tokens with character offsets.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
list[PhraseMatch]: Deduplicated unhydrated phrase matches with token and character spans.
|
||||
"""
|
||||
matches: list[PhraseMatch] = []
|
||||
seen: set[tuple[int | None, str, int, int]] = set()
|
||||
token_texts = [token.text for token in tokens_]
|
||||
for phrase_norm, start, end in generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens):
|
||||
phrase_ids = lookup.norm_to_phrase_ids.get(phrase_norm, ())
|
||||
alias_ids = lookup.alias_to_phrase_ids.get(phrase_norm, ())
|
||||
for phrase_id in (*phrase_ids, *alias_ids):
|
||||
key = (phrase_id, phrase_norm, start, end)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
matches.append(
|
||||
PhraseMatch(
|
||||
phrase_norm=phrase_norm,
|
||||
phrase_id=phrase_id,
|
||||
start_token=start,
|
||||
end_token=end,
|
||||
token_count=end - start,
|
||||
start_char=tokens_[start].start_char,
|
||||
end_char=tokens_[end - 1].end_char,
|
||||
)
|
||||
)
|
||||
return matches
|
||||
|
||||
|
||||
async def hydrate_matches(session: AsyncSession, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]:
|
||||
"""Fetch protected phrase metadata for raw phrase matches.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
matches (Sequence[PhraseMatch]): Unhydrated matches to enrich.
|
||||
|
||||
Returns:
|
||||
list[HydratedPhraseMatch]: Matches with protected-phrase metadata attached.
|
||||
"""
|
||||
if not matches:
|
||||
return []
|
||||
|
||||
phrase_ids = sorted({match.phrase_id for match in matches if match.phrase_id is not None})
|
||||
if not phrase_ids:
|
||||
return []
|
||||
|
||||
rows = {
|
||||
row.id: row
|
||||
for row in await session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids)))
|
||||
}
|
||||
hydrated: list[HydratedPhraseMatch] = []
|
||||
for match in matches:
|
||||
if match.phrase_id is None:
|
||||
continue
|
||||
phrase = rows.get(match.phrase_id)
|
||||
if phrase is None:
|
||||
continue
|
||||
hydrated.append(
|
||||
HydratedPhraseMatch(
|
||||
phrase_id=phrase.id,
|
||||
matched_norm=match.phrase_norm,
|
||||
phrase_text=phrase.phrase_text,
|
||||
phrase_norm=phrase.phrase_norm,
|
||||
canonical_id=phrase.canonical_id,
|
||||
phrase_type=phrase.phrase_type,
|
||||
token_count=match.token_count,
|
||||
confidence=phrase.confidence,
|
||||
importance=phrase.importance,
|
||||
allow_nested=phrase.allow_nested,
|
||||
suppress_children=phrase.suppress_children,
|
||||
start_token=match.start_token,
|
||||
end_token=match.end_token,
|
||||
start_char=match.start_char,
|
||||
end_char=match.end_char,
|
||||
book_id=phrase.book_id,
|
||||
series_id=phrase.series_id,
|
||||
)
|
||||
)
|
||||
return hydrated
|
||||
|
||||
|
||||
def overlaps(first: HydratedPhraseMatch, second: HydratedPhraseMatch) -> bool:
|
||||
"""Return whether two token spans overlap.
|
||||
|
||||
Args:
|
||||
first (HydratedPhraseMatch): First match to compare.
|
||||
second (HydratedPhraseMatch): Second match to compare.
|
||||
|
||||
Returns:
|
||||
bool: True when the two token spans share at least one token position.
|
||||
"""
|
||||
return not (first.end_token <= second.start_token or first.start_token >= second.end_token)
|
||||
|
||||
|
||||
def is_inside(child: HydratedPhraseMatch, parent: HydratedPhraseMatch) -> bool:
|
||||
"""Return whether one token span is strictly inside another.
|
||||
|
||||
Args:
|
||||
child (HydratedPhraseMatch): Candidate nested match.
|
||||
parent (HydratedPhraseMatch): Candidate enclosing match.
|
||||
|
||||
Returns:
|
||||
bool: True when ``child`` lies within ``parent`` and is not the same span.
|
||||
"""
|
||||
return (
|
||||
child.start_token >= parent.start_token
|
||||
and child.end_token <= parent.end_token
|
||||
and (child.start_token, child.end_token, child.phrase_id)
|
||||
!= (parent.start_token, parent.end_token, parent.phrase_id)
|
||||
)
|
||||
|
||||
|
||||
def rank_match(match: HydratedPhraseMatch) -> tuple[float, float, int]:
|
||||
"""Rank phrase matches by importance, confidence, then token count.
|
||||
|
||||
Args:
|
||||
match (HydratedPhraseMatch): Match to build a sort key for.
|
||||
|
||||
Returns:
|
||||
tuple[float, float, int]: A comparable key of importance, confidence, and token count.
|
||||
"""
|
||||
return (match.importance, match.confidence, match.token_count)
|
||||
|
||||
|
||||
def should_suppress(candidate: HydratedPhraseMatch, kept: HydratedPhraseMatch) -> bool:
|
||||
"""Return whether an already-kept match should suppress a candidate.
|
||||
|
||||
Args:
|
||||
candidate (HydratedPhraseMatch): Match being considered for keeping.
|
||||
kept (HydratedPhraseMatch): Match already kept that may suppress the candidate.
|
||||
|
||||
Returns:
|
||||
bool: True when the candidate should be dropped in favor of the kept match.
|
||||
"""
|
||||
if not overlaps(candidate, kept):
|
||||
return False
|
||||
if candidate.canonical_id == kept.canonical_id:
|
||||
return rank_match(kept) >= rank_match(candidate)
|
||||
if is_inside(candidate, kept) and kept.suppress_children and not candidate.allow_nested:
|
||||
return True
|
||||
return not candidate.allow_nested and rank_match(kept) > rank_match(candidate)
|
||||
|
||||
|
||||
def resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhraseMatch]:
|
||||
"""Resolve overlapping phrase matches without relying only on longest match.
|
||||
|
||||
Args:
|
||||
matches (Sequence[HydratedPhraseMatch]): Hydrated matches that may overlap.
|
||||
|
||||
Returns:
|
||||
list[HydratedPhraseMatch]: The kept, non-suppressed matches.
|
||||
"""
|
||||
sorted_matches = sorted(
|
||||
matches,
|
||||
key=lambda match: (match.start_token, -match.token_count, -match.importance, -match.confidence),
|
||||
)
|
||||
kept: list[HydratedPhraseMatch] = []
|
||||
for candidate in sorted_matches:
|
||||
if any(should_suppress(candidate, existing) for existing in kept):
|
||||
continue
|
||||
kept.append(candidate)
|
||||
return kept
|
||||
|
||||
|
||||
async def detect_protected_phrases_for_query(
|
||||
session: AsyncSession,
|
||||
query_text: str,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
lookup: PhraseLookup | None = None,
|
||||
book_id: int | None = None,
|
||||
series_id: int | None = None,
|
||||
) -> list[HydratedPhraseMatch]:
|
||||
"""Run the full online protected-phrase query-detection pipeline.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
query_text (str): User query text to detect phrases in.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
|
||||
book_id (int | None): Optional book scope for lookup loading.
|
||||
series_id (int | None): Optional series scope for lookup loading.
|
||||
|
||||
Returns:
|
||||
list[HydratedPhraseMatch]: Hydrated, overlap-resolved phrase matches for the query.
|
||||
"""
|
||||
active_lookup = (
|
||||
lookup
|
||||
if lookup is not None
|
||||
else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
|
||||
)
|
||||
return resolve_overlaps(await hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup)))
|
||||
|
||||
|
||||
async def index_chunk_phrase_mentions_for_book(
|
||||
session: AsyncSession,
|
||||
book_id: int,
|
||||
config: EbookSearchConfig,
|
||||
*,
|
||||
series_id: int | None = None,
|
||||
lookup: PhraseLookup | None = None,
|
||||
) -> int:
|
||||
"""Rebuild chunk phrase mentions for all chunks in one book.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
book_id (int): Book whose chunk mentions are rebuilt.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
series_id (int | None): Optional series scope for lookup loading.
|
||||
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
|
||||
|
||||
Returns:
|
||||
int: Total number of chunk phrase mentions indexed for the book.
|
||||
"""
|
||||
active_lookup = (
|
||||
lookup
|
||||
if lookup is not None
|
||||
else await load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
|
||||
)
|
||||
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
|
||||
chunks = await session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id))
|
||||
count = 0
|
||||
for chunk in chunks:
|
||||
count += await index_chunk_phrase_mentions(session, chunk, lookup=active_lookup)
|
||||
await session.flush()
|
||||
logger.info(f"ebook_chunk_phrase_mentions_indexed {book_id=} {count=}")
|
||||
return count
|
||||
|
||||
|
||||
async def index_chunk_phrase_mentions(session: AsyncSession, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
|
||||
"""Store protected phrase mentions for one chunk.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
chunk (EbookChunk): Chunk whose text is scanned for phrase mentions.
|
||||
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
|
||||
|
||||
Returns:
|
||||
int: Number of phrase mentions stored for the chunk.
|
||||
"""
|
||||
raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup)
|
||||
hydrated = resolve_overlaps(await hydrate_matches(session, raw_matches))
|
||||
for match in hydrated:
|
||||
session.add(
|
||||
EbookChunkPhraseMention(
|
||||
chunk_id=chunk.id,
|
||||
phrase_id=match.phrase_id,
|
||||
book_id=match.book_id if match.book_id is not None else chunk.source_id,
|
||||
series_id=match.series_id,
|
||||
start_char=match.start_char if match.start_char is not None else 0,
|
||||
end_char=match.end_char,
|
||||
)
|
||||
)
|
||||
return len(hydrated)
|
||||
|
||||
|
||||
async def phrase_hits_for_chunks(
|
||||
session: AsyncSession,
|
||||
*,
|
||||
chunk_ids: Sequence[int],
|
||||
phrase_ids: Sequence[int],
|
||||
) -> dict[int, tuple[ChunkPhraseHit, ...]]:
|
||||
"""Return matched protected phrases with mention counts by chunk id using indexed chunk mentions.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
chunk_ids (Sequence[int]): Chunk ids to look up mentions for.
|
||||
phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to.
|
||||
|
||||
Returns:
|
||||
dict[int, tuple[ChunkPhraseHit, ...]]: Phrase hits per chunk id, ordered by mention count.
|
||||
"""
|
||||
if not chunk_ids or not phrase_ids:
|
||||
return {}
|
||||
|
||||
mention_count = func.count(EbookChunkPhraseMention.phrase_id).label("mention_count")
|
||||
statement = (
|
||||
select(
|
||||
EbookChunkPhraseMention.chunk_id,
|
||||
EbookProtectedPhrase.id.label("phrase_id"),
|
||||
EbookProtectedPhrase.phrase_text,
|
||||
mention_count,
|
||||
)
|
||||
.join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookChunkPhraseMention.phrase_id)
|
||||
.where(
|
||||
EbookChunkPhraseMention.chunk_id.in_(chunk_ids),
|
||||
EbookChunkPhraseMention.phrase_id.in_(phrase_ids),
|
||||
)
|
||||
.group_by(EbookChunkPhraseMention.chunk_id, EbookProtectedPhrase.id, EbookProtectedPhrase.phrase_text)
|
||||
.order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text)
|
||||
)
|
||||
hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list)
|
||||
for row in await session.execute(statement):
|
||||
hits[int(row.chunk_id)].append(
|
||||
ChunkPhraseHit(
|
||||
phrase_id=int(row.phrase_id),
|
||||
phrase_text=str(row.phrase_text),
|
||||
mention_count=int(row.mention_count),
|
||||
)
|
||||
)
|
||||
return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()}
|
||||
|
||||
|
||||
async def phrase_hit_counts_for_chunks(
|
||||
session: AsyncSession,
|
||||
*,
|
||||
chunk_ids: Sequence[int],
|
||||
phrase_ids: Sequence[int],
|
||||
) -> dict[int, int]:
|
||||
"""Return phrase-hit counts by chunk id using indexed chunk mentions.
|
||||
|
||||
Args:
|
||||
session (AsyncSession): Active database session.
|
||||
chunk_ids (Sequence[int]): Chunk ids to count mentions for.
|
||||
phrase_ids (Sequence[int]): Protected phrase ids to restrict the counts to.
|
||||
|
||||
Returns:
|
||||
dict[int, int]: Total mention count per chunk id.
|
||||
"""
|
||||
hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
|
||||
return {chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in hits.items()}
|
||||
|
||||
@@ -8,6 +8,8 @@ from typing import TYPE_CHECKING
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Mapping
|
||||
|
||||
from python.orm.richie import EbookProtectedPhrase
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PhraseCandidate:
|
||||
@@ -71,47 +73,24 @@ class LLMJudgment:
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PhraseLookup:
|
||||
"""In-memory lookup maps used for constant-time phrase-window checks.
|
||||
"""In-memory phrase metadata used for constant-time text-window checks.
|
||||
|
||||
Attributes:
|
||||
norm_to_phrase_ids (Mapping[str, tuple[int, ...]]): Normalized phrase to protected phrase ids.
|
||||
alias_to_phrase_ids (Mapping[str, tuple[int, ...]]): Normalized alias to protected phrase ids.
|
||||
phrase_ids_by_norm (Mapping[str, tuple[int, ...]]): Canonical and alias norms to phrase ids.
|
||||
phrases_by_id (Mapping[int, EbookProtectedPhrase]): Protected phrase metadata by id.
|
||||
min_tokens (int): Smallest token-window size to test.
|
||||
max_tokens (int): Largest token-window size to test.
|
||||
"""
|
||||
|
||||
norm_to_phrase_ids: Mapping[str, tuple[int, ...]]
|
||||
alias_to_phrase_ids: Mapping[str, tuple[int, ...]]
|
||||
phrase_ids_by_norm: Mapping[str, tuple[int, ...]]
|
||||
phrases_by_id: Mapping[int, EbookProtectedPhrase]
|
||||
min_tokens: int
|
||||
max_tokens: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PhraseMatch:
|
||||
"""An unhydrated query or chunk phrase match.
|
||||
|
||||
Attributes:
|
||||
phrase_norm (str): Normalized text of the matched window.
|
||||
start_token (int): Index of the first matched token.
|
||||
end_token (int): Index one past the last matched token.
|
||||
token_count (int): Number of tokens in the match.
|
||||
phrase_id (int | None): Matched protected phrase id when known.
|
||||
start_char (int | None): Start character offset in the source text.
|
||||
end_char (int | None): End character offset in the source text.
|
||||
"""
|
||||
|
||||
phrase_norm: str
|
||||
start_token: int
|
||||
end_token: int
|
||||
token_count: int
|
||||
phrase_id: int | None = None
|
||||
start_char: int | None = None
|
||||
end_char: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class HydratedPhraseMatch:
|
||||
"""A phrase match with protected-phrase metadata attached.
|
||||
"""A detected phrase match with protected-phrase metadata attached.
|
||||
|
||||
Attributes:
|
||||
phrase_id (int): Protected phrase id.
|
||||
@@ -152,21 +131,6 @@ class HydratedPhraseMatch:
|
||||
series_id: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ChunkPhraseHit:
|
||||
"""One protected phrase with its mention count inside one retrieved chunk.
|
||||
|
||||
Attributes:
|
||||
phrase_id (int): Protected phrase id.
|
||||
phrase_text (str): Display text of the protected phrase.
|
||||
mention_count (int): Indexed mentions of the phrase in the chunk.
|
||||
"""
|
||||
|
||||
phrase_id: int
|
||||
phrase_text: str
|
||||
mention_count: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PhraseCandidateGenerationResult:
|
||||
"""Summary of candidate phrase extraction for indexed books.
|
||||
|
||||
@@ -19,6 +19,7 @@ from python.ebook_search.bm25_corpus import (
|
||||
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,
|
||||
@@ -40,7 +41,7 @@ if TYPE_CHECKING:
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
from python.ebook_search.protected_phrases.models import HydratedPhraseMatch
|
||||
from python.ebook_search.protected_phrases.models import PhraseMatch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -74,7 +75,7 @@ class SearchResponse:
|
||||
results: list[SearchResult]
|
||||
rank_label: str
|
||||
timings: tuple[RuntimeStep, ...] = ()
|
||||
phrase_matches: tuple[HydratedPhraseMatch, ...] = ()
|
||||
phrase_matches: tuple[PhraseMatch, ...] = ()
|
||||
|
||||
@property
|
||||
def total_runtime_ms(self) -> float:
|
||||
@@ -88,7 +89,7 @@ class RetrievalResponse:
|
||||
|
||||
vector_results: list[SearchResult]
|
||||
lexical_results: list[SearchResult]
|
||||
phrase_matches: list[HydratedPhraseMatch]
|
||||
phrase_matches: list[PhraseMatch]
|
||||
timings: tuple[RuntimeStep, ...]
|
||||
|
||||
|
||||
@@ -151,18 +152,17 @@ async 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 []
|
||||
|
||||
|
||||
async def query_phrase_matches(
|
||||
engine: AsyncEngine,
|
||||
query: str,
|
||||
config: EbookSearchConfig,
|
||||
) -> list[HydratedPhraseMatch]:
|
||||
*,
|
||||
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)
|
||||
@@ -180,7 +180,7 @@ def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchRes
|
||||
async def apply_phrase_mention_boosts(
|
||||
engine: AsyncEngine,
|
||||
candidates: list[SearchResult],
|
||||
phrase_matches: Sequence[HydratedPhraseMatch],
|
||||
phrase_matches: Sequence[PhraseMatch],
|
||||
phrase_hit_boost: float,
|
||||
) -> list[SearchResult]:
|
||||
"""Boost retrieved chunks that have indexed mentions for detected protected phrases."""
|
||||
@@ -242,23 +242,21 @@ async def parallel_retrieval(
|
||||
|
||||
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.
|
||||
it joins the gather as a third task and returns immediately when phrase matching is disabled.
|
||||
"""
|
||||
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(
|
||||
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),
|
||||
),
|
||||
)
|
||||
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)} "
|
||||
@@ -334,13 +332,7 @@ async def vector_candidates(
|
||||
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"),
|
||||
*CHUNK_RECORD_COLUMNS,
|
||||
score,
|
||||
)
|
||||
.select_from(embedding_table)
|
||||
|
||||
@@ -25,7 +25,7 @@ from python.ebook_search.bm25_corpus import (
|
||||
score_bm25_corpus,
|
||||
write_bm25_corpus,
|
||||
)
|
||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config, normalize_embedding_model
|
||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_config
|
||||
from python.ebook_search.embeddings import MODEL_DIMENSIONS, ensure_embedding_models
|
||||
from python.ebook_search.ingest import chunk_text, find_existing_source
|
||||
from python.ebook_search.search import (
|
||||
@@ -452,24 +452,25 @@ def test_1024_embedding_table_has_cosine_hnsw_index() -> None:
|
||||
|
||||
def test_embedding_model_aliases_normalize_to_provider_names(mocker: MockerFixture) -> None:
|
||||
mocker.patch.dict(environ, {}, clear=False)
|
||||
environ.pop("EBOOK_SEARCH_EMBEDDING_MODEL", None)
|
||||
|
||||
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
|
||||
assert load_config().embedding_model == "qwen3-embedding-0.6b"
|
||||
|
||||
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-0.6b"
|
||||
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
|
||||
assert load_config().embedding_model == "qwen3-embedding-0.6b"
|
||||
|
||||
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen3-Embedding-0.6B"
|
||||
assert normalize_embedding_model() == "qwen3-embedding-0.6b"
|
||||
assert load_config().embedding_model == "qwen3-embedding-0.6b"
|
||||
|
||||
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "Qwen/Qwen3-Embedding-4B"
|
||||
|
||||
assert normalize_embedding_model() == "qwen3-embedding-4b"
|
||||
assert load_config().embedding_model == "qwen3-embedding-4b"
|
||||
|
||||
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding:8b"
|
||||
assert normalize_embedding_model() == "qwen3-embedding-8b"
|
||||
assert load_config().embedding_model == "qwen3-embedding-8b"
|
||||
|
||||
environ["EBOOK_SEARCH_EMBEDDING_MODEL"] = "qwen3-embedding-8b"
|
||||
assert normalize_embedding_model() == "qwen3-embedding-8b"
|
||||
assert load_config().embedding_model == "qwen3-embedding-8b"
|
||||
|
||||
|
||||
def test_answer_generation_is_enabled_by_default(mocker: MockerFixture) -> None:
|
||||
|
||||
@@ -10,6 +10,7 @@ import pytest
|
||||
from python.ebook_search.answer import answer_query
|
||||
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
||||
from python.ebook_search.embeddings import embed_texts
|
||||
from python.ebook_search.llm_interface import check_chat_endpoint, check_embedding_endpoint
|
||||
from python.ebook_search.search import SearchResult
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -23,6 +24,44 @@ def make_async_client(mocker: MockerFixture, fake_post) -> httpx.AsyncClient:
|
||||
return client
|
||||
|
||||
|
||||
async def test_model_endpoint_checks_share_http_probe(mocker: MockerFixture) -> None:
|
||||
client = mocker.MagicMock(spec=httpx.AsyncClient)
|
||||
response = mocker.MagicMock(spec=httpx.Response)
|
||||
client.get = mocker.AsyncMock(return_value=response)
|
||||
config = EbookSearchConfig(
|
||||
rerank=RerankConfig(enabled=False),
|
||||
embedding_base_url="https://embedding.example/v1/",
|
||||
vllm_base_url="https://chat.example/v1/",
|
||||
vllm_api_key="secret",
|
||||
)
|
||||
|
||||
assert await check_embedding_endpoint(client, config, timeout_seconds=2.0)
|
||||
assert await check_chat_endpoint(client, config, timeout_seconds=3.0)
|
||||
assert client.get.await_args_list == [
|
||||
mocker.call("https://embedding.example/v1/models", headers={}, timeout=2.0),
|
||||
mocker.call(
|
||||
"https://chat.example/v1/models",
|
||||
headers={"Authorization": "Bearer secret"},
|
||||
timeout=3.0,
|
||||
),
|
||||
]
|
||||
assert response.raise_for_status.call_count == 2
|
||||
|
||||
|
||||
async def test_model_endpoint_checks_report_http_failures(mocker: MockerFixture) -> None:
|
||||
client = mocker.MagicMock(spec=httpx.AsyncClient)
|
||||
client.get = mocker.AsyncMock(
|
||||
side_effect=[
|
||||
httpx.ConnectError("embedding offline"),
|
||||
httpx.ConnectError("chat offline"),
|
||||
]
|
||||
)
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
|
||||
assert not await check_embedding_endpoint(client, config)
|
||||
assert not await check_chat_endpoint(client, config)
|
||||
|
||||
|
||||
async def test_answer_query_uses_httpx_chat_completions(mocker: MockerFixture) -> None:
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
|
||||
@@ -7,7 +7,7 @@ from datetime import UTC, datetime
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy import event, select
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
@@ -33,15 +33,12 @@ from python.ebook_search.protected_phrases.matching import (
|
||||
detect_protected_phrases_for_query,
|
||||
index_chunk_phrase_mentions,
|
||||
load_phrase_lookup,
|
||||
phrase_hit_counts_for_chunks,
|
||||
phrase_hits_for_chunks,
|
||||
resolve_overlaps,
|
||||
)
|
||||
from python.ebook_search.protected_phrases.models import (
|
||||
ChunkPhraseHit,
|
||||
HydratedPhraseMatch,
|
||||
LLMJudgment,
|
||||
PhraseCandidate,
|
||||
PhraseMatch,
|
||||
)
|
||||
from python.ebook_search.protected_phrases.store import book_ids_pending_first_judgment, corpus_phrase_stats
|
||||
from python.ebook_search.protected_phrases.text_normalization import normalize_text, tokenize
|
||||
@@ -187,26 +184,89 @@ def test_score_candidate_weights_metadata_source(config: EbookSearchConfig) -> N
|
||||
assert score_candidate(metadata_only, config) == 2.0 + 0.5
|
||||
|
||||
|
||||
async def test_detect_protected_phrases_hydrates_alias_matches(
|
||||
async def test_load_phrase_lookup_loads_phrases_and_aliases_with_one_query(
|
||||
engine: AsyncEngine,
|
||||
session: AsyncSession,
|
||||
) -> None:
|
||||
"""One joined query should load phrases with zero, one, or many aliases."""
|
||||
source = await add_source(session)
|
||||
aliased_phrase = await add_phrase(session, source.id, phrase_text="lock in", phrase_norm="lock in")
|
||||
plain_phrase = await add_phrase(session, source.id, phrase_text="mage king", phrase_norm="mage king")
|
||||
session.add_all(
|
||||
[
|
||||
EbookPhraseAlias(phrase_id=aliased_phrase.id, alias_text="locked in", alias_norm="locked in"),
|
||||
EbookPhraseAlias(
|
||||
phrase_id=aliased_phrase.id,
|
||||
alias_text="locked completely in",
|
||||
alias_norm="locked completely in",
|
||||
),
|
||||
]
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
statements: list[str] = []
|
||||
|
||||
def record_statement(*args: object) -> None:
|
||||
statements.append(str(args[2]))
|
||||
|
||||
event.listen(engine.sync_engine, "before_cursor_execute", record_statement)
|
||||
try:
|
||||
lookup = await load_phrase_lookup(session, EbookSearchConfig(phrase_max_tokens=2), book_id=source.id)
|
||||
finally:
|
||||
event.remove(engine.sync_engine, "before_cursor_execute", record_statement)
|
||||
|
||||
assert lookup.phrase_ids_by_norm == {
|
||||
"lock in": (aliased_phrase.id,),
|
||||
"locked completely in": (aliased_phrase.id,),
|
||||
"locked in": (aliased_phrase.id,),
|
||||
"mage king": (plain_phrase.id,),
|
||||
}
|
||||
assert lookup.phrases_by_id == {aliased_phrase.id: aliased_phrase, plain_phrase.id: plain_phrase}
|
||||
assert lookup.max_tokens == 3
|
||||
assert len(statements) == 1
|
||||
|
||||
|
||||
async def test_detect_protected_phrases_queries_canonical_and_alias_matches_once(
|
||||
engine: AsyncEngine,
|
||||
session: AsyncSession,
|
||||
config: EbookSearchConfig,
|
||||
) -> None:
|
||||
"""Query detection should use RAM aliases and hydrate phrase metadata from the DB."""
|
||||
"""Query detection should hydrate repeated canonical and alias matches with one statement."""
|
||||
source = await add_source(session)
|
||||
phrase = await add_phrase(session, source.id, phrase_text="lock in", phrase_norm="lock in")
|
||||
await add_phrase(session, source.id, phrase_text="mage king", phrase_norm="mage king")
|
||||
session.add(EbookPhraseAlias(phrase_id=phrase.id, alias_text="locked in", alias_norm="locked in"))
|
||||
await session.commit()
|
||||
|
||||
matches = await detect_protected_phrases_for_query(session, "what is locked-in", config)
|
||||
statements: list[str] = []
|
||||
|
||||
assert [(match.phrase_text, match.canonical_id, match.phrase_type) for match in matches] == [
|
||||
("lock in", "condition:lock_in", "fictional_condition")
|
||||
def record_statement(*args: object) -> None:
|
||||
statements.append(str(args[2]))
|
||||
|
||||
event.listen(engine.sync_engine, "before_cursor_execute", record_statement)
|
||||
try:
|
||||
matches = await detect_protected_phrases_for_query(
|
||||
session,
|
||||
"lock-in then locked-in and lock-in",
|
||||
config,
|
||||
)
|
||||
empty_matches = await detect_protected_phrases_for_query(session, "", config)
|
||||
finally:
|
||||
event.remove(engine.sync_engine, "before_cursor_execute", record_statement)
|
||||
|
||||
assert [(match.phrase_text, match.matched_norm, match.start_token, match.end_token) for match in matches] == [
|
||||
("lock in", "lock in", 0, 2),
|
||||
("lock in", "locked in", 3, 5),
|
||||
("lock in", "lock in", 6, 8),
|
||||
]
|
||||
assert empty_matches == []
|
||||
assert len(statements) == 1
|
||||
assert " UNION " in statements[0]
|
||||
|
||||
|
||||
def test_resolve_overlaps_keeps_independent_nested_phrases() -> None:
|
||||
"""Overlap resolution should keep useful nested concepts when metadata permits it."""
|
||||
child = hydrated_match(
|
||||
child = phrase_match(
|
||||
phrase_id=1,
|
||||
phrase_text="mage king",
|
||||
canonical_id="title:mage_king",
|
||||
@@ -214,7 +274,7 @@ def test_resolve_overlaps_keeps_independent_nested_phrases() -> None:
|
||||
end_token=5,
|
||||
allow_nested=True,
|
||||
)
|
||||
parent = hydrated_match(
|
||||
parent = phrase_match(
|
||||
phrase_id=2,
|
||||
phrase_text="mage king of mars",
|
||||
canonical_id="entity:mage_king_of_mars",
|
||||
@@ -231,7 +291,7 @@ def test_resolve_overlaps_keeps_independent_nested_phrases() -> None:
|
||||
|
||||
def test_resolve_overlaps_suppresses_weaker_same_canonical_match() -> None:
|
||||
"""Same-canonical overlaps should keep the stronger evidence."""
|
||||
weak = hydrated_match(
|
||||
weak = phrase_match(
|
||||
phrase_id=1,
|
||||
phrase_text="lock",
|
||||
canonical_id="condition:lock_in",
|
||||
@@ -239,7 +299,7 @@ def test_resolve_overlaps_suppresses_weaker_same_canonical_match() -> None:
|
||||
end_token=3,
|
||||
importance=0.2,
|
||||
)
|
||||
strong = hydrated_match(
|
||||
strong = phrase_match(
|
||||
phrase_id=2,
|
||||
phrase_text="lock in",
|
||||
canonical_id="condition:lock_in",
|
||||
@@ -274,19 +334,25 @@ async def test_index_chunk_phrase_mentions_uses_normalized_window_lookup(
|
||||
await session.commit()
|
||||
lookup = await load_phrase_lookup(session, config, book_id=source.id)
|
||||
|
||||
count = await index_chunk_phrase_mentions(session, chunk, lookup=lookup)
|
||||
statements: list[str] = []
|
||||
|
||||
def record_statement(*args: object) -> None:
|
||||
statements.append(str(args[2]))
|
||||
|
||||
event.listen(session.bind.sync_engine, "before_cursor_execute", record_statement)
|
||||
try:
|
||||
count = index_chunk_phrase_mentions(session, chunk, lookup=lookup)
|
||||
finally:
|
||||
event.remove(session.bind.sync_engine, "before_cursor_execute", record_statement)
|
||||
await session.commit()
|
||||
|
||||
assert statements == []
|
||||
mention = await session.scalar(select(EbookChunkPhraseMention))
|
||||
assert count == 1
|
||||
assert mention is not None
|
||||
assert mention.chunk_id == chunk.id
|
||||
assert mention.phrase_id == phrase.id
|
||||
assert chunk.text[mention.start_char : mention.end_char] == "lock-in"
|
||||
assert await phrase_hit_counts_for_chunks(session, chunk_ids=[chunk.id], phrase_ids=[phrase.id]) == {chunk.id: 1}
|
||||
assert await phrase_hits_for_chunks(session, chunk_ids=[chunk.id], phrase_ids=[phrase.id]) == {
|
||||
chunk.id: (ChunkPhraseHit(phrase_id=phrase.id, phrase_text="lock in", mention_count=1),)
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("worker_pool")
|
||||
@@ -1060,7 +1126,7 @@ async def add_phrase(
|
||||
return phrase
|
||||
|
||||
|
||||
def hydrated_match(
|
||||
def phrase_match(
|
||||
*,
|
||||
phrase_id: int,
|
||||
phrase_text: str,
|
||||
@@ -1070,9 +1136,9 @@ def hydrated_match(
|
||||
importance: float = 0.8,
|
||||
allow_nested: bool = False,
|
||||
suppress_children: bool = True,
|
||||
) -> HydratedPhraseMatch:
|
||||
"""Build a hydrated match for overlap tests."""
|
||||
return HydratedPhraseMatch(
|
||||
) -> PhraseMatch:
|
||||
"""Build a metadata-backed phrase match for overlap tests."""
|
||||
return PhraseMatch(
|
||||
phrase_id=phrase_id,
|
||||
matched_norm=phrase_text,
|
||||
phrase_text=phrase_text,
|
||||
|
||||
@@ -62,8 +62,9 @@ async def test_search_ebooks_runs_phrase_detection_in_parallel_with_retrieval(mo
|
||||
assert await asyncio.to_thread(phrase_started.wait, 2)
|
||||
return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]
|
||||
|
||||
async def fake_query_phrase_matches(_engine, _query, _config):
|
||||
async def fake_query_phrase_matches(_engine, _query, _config, *, phrase_matching):
|
||||
"""Record that phrase detection started and return no matches."""
|
||||
assert phrase_matching is True
|
||||
phrase_started.set()
|
||||
return []
|
||||
|
||||
@@ -90,7 +91,7 @@ async def test_search_ebooks_skips_phrase_matching_when_disabled(mocker: MockerF
|
||||
return_value=[SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)],
|
||||
)
|
||||
mocker.patch("python.ebook_search.search.bm25_candidates", return_value=[])
|
||||
detect_mock = mocker.patch("python.ebook_search.search.query_phrase_matches")
|
||||
detect_mock = mocker.patch("python.ebook_search.search.detect_protected_phrases_for_query")
|
||||
boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
|
||||
|
||||
@@ -115,7 +116,7 @@ async def test_search_ebooks_ignores_phrase_matching_when_config_disabled(mocker
|
||||
return_value=[SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)],
|
||||
)
|
||||
mocker.patch("python.ebook_search.search.bm25_candidates", return_value=[])
|
||||
detect_mock = mocker.patch("python.ebook_search.search.query_phrase_matches")
|
||||
detect_mock = mocker.patch("python.ebook_search.search.detect_protected_phrases_for_query")
|
||||
boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
|
||||
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), phrase_matching_enabled=False)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user