Add models and database persistence for protected phrase extraction
- Introduced dataclasses for phrase candidates, judgments, and matches in `models.py`. - Implemented database operations for candidate and protected phrases in `store.py`, including loading, saving, and deleting phrases. - Enhanced text normalization functions in `text_normalization.py` with detailed docstrings. - Refactored search functionality to utilize new models and methods for detecting protected phrases.
This commit is contained in:
@@ -0,0 +1,481 @@
|
||||
"""Book-level orchestration for LLM judging and promotion of candidate phrases."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
from time import perf_counter
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import select
|
||||
|
||||
from python.ebook_search.llm_interface import request_chat_completion
|
||||
from python.ebook_search.protected_phrases.extraction import (
|
||||
candidate_source_names,
|
||||
get_sample_contexts,
|
||||
is_most_common_word_phrase,
|
||||
minimum_candidate_raw_count,
|
||||
)
|
||||
from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book
|
||||
from python.ebook_search.protected_phrases.models import BookJudgmentResult, LLMJudgment, PhraseJudgmentBackfillResult
|
||||
from python.ebook_search.protected_phrases.store import (
|
||||
count_protected_phrases,
|
||||
count_unjudged_candidates,
|
||||
load_book_text,
|
||||
load_candidates_for_judgment,
|
||||
phrase_candidate_from_row,
|
||||
save_candidate_to_db,
|
||||
upsert_protected_phrase,
|
||||
)
|
||||
from python.ebook_search.protected_phrases.text_normalization import normalize_text
|
||||
from python.orm.richie import EbookSource
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
from python.ebook_search.protected_phrases.models import PhraseCandidate
|
||||
from python.orm.richie import EbookCandidatePhrase, EbookProtectedPhrase
|
||||
|
||||
JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)
|
||||
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def judge_candidate_phrases_for_books(
|
||||
session: Session,
|
||||
config: EbookSearchConfig,
|
||||
) -> PhraseJudgmentBackfillResult:
|
||||
"""Judge stored candidate phrases and promote accepted phrases for indexed books.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
PhraseJudgmentBackfillResult: Per-corpus counts of books judged, failures, candidates,
|
||||
protected phrases, and mentions.
|
||||
"""
|
||||
source_ids = session.scalars(select(EbookSource.id).order_by(EbookSource.id)).all()
|
||||
books_seen = len(source_ids)
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_start books_seen=%s llm_candidates_per_book=%s "
|
||||
"target_protected_per_book=%s confidence_threshold=%.2f min_tokens=%s min_uses=%s",
|
||||
books_seen,
|
||||
config.protected_phrase_llm_candidates_per_book,
|
||||
config.phrase_target_protected_per_book,
|
||||
config.protected_phrase_confidence_threshold,
|
||||
config.phrase_min_tokens,
|
||||
minimum_candidate_raw_count(config),
|
||||
)
|
||||
|
||||
outcomes = [judge_book_for_backfill(session, source_id, config) for source_id in source_ids]
|
||||
|
||||
result = PhraseJudgmentBackfillResult(
|
||||
books_seen=books_seen,
|
||||
books_judged=sum(1 for outcome in outcomes if outcome.committed),
|
||||
books_failed=sum(1 for outcome in outcomes if outcome.failed),
|
||||
candidates_judged=sum(outcome.judged for outcome in outcomes),
|
||||
protected_phrases=sum(outcome.protected for outcome in outcomes),
|
||||
phrase_mentions=sum(outcome.mentions for outcome in outcomes),
|
||||
)
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_complete books_seen=%s books_judged=%s books_failed=%s "
|
||||
"candidates_judged=%s protected=%s mentions=%s",
|
||||
result.books_seen,
|
||||
result.books_judged,
|
||||
result.books_failed,
|
||||
result.candidates_judged,
|
||||
result.protected_phrases,
|
||||
result.phrase_mentions,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def judge_book_for_backfill(
|
||||
session: Session,
|
||||
source_id: int,
|
||||
config: EbookSearchConfig,
|
||||
) -> BookJudgmentResult:
|
||||
"""Judge one book's candidates and return its outcome.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
source_id (int): Book to judge candidates for.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
BookJudgmentResult: The book's judgment outcome, or an empty result when nothing was unjudged.
|
||||
"""
|
||||
unjudged_count = count_unjudged_candidates(session, source_id, config)
|
||||
if not unjudged_count:
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s",
|
||||
source_id,
|
||||
)
|
||||
return BookJudgmentResult()
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_book_start source_id=%s unjudged=%s",
|
||||
source_id,
|
||||
unjudged_count,
|
||||
)
|
||||
return run_book_judgment(session, source_id, unjudged_count, config)
|
||||
|
||||
|
||||
def run_book_judgment(
|
||||
session: Session,
|
||||
source_id: int,
|
||||
unjudged_count: int,
|
||||
config: EbookSearchConfig,
|
||||
) -> BookJudgmentResult:
|
||||
"""Judge and index one book's candidates, managing its own transaction.
|
||||
|
||||
Commits on success and rolls back on error, returning a :class:`BookJudgmentResult`
|
||||
that describes the outcome rather than raising it to the caller.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
source_id (int): Book to judge candidates for.
|
||||
unjudged_count (int): Storable unjudged candidates counted before judging, for logging.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
BookJudgmentResult: Judged, protected, and mention counts with commit and failure flags.
|
||||
"""
|
||||
book_started_at = perf_counter()
|
||||
try:
|
||||
book_text = load_book_text(session, source_id)
|
||||
if not book_text:
|
||||
logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id)
|
||||
return BookJudgmentResult()
|
||||
|
||||
judged, protected = judge_candidate_phrases_for_book(
|
||||
session,
|
||||
source_id,
|
||||
series_id=None,
|
||||
normalized_book_text=normalize_text(book_text),
|
||||
config=config,
|
||||
)
|
||||
if judged == 0:
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_book_skip_no_judgments source_id=%s unjudged=%s",
|
||||
source_id,
|
||||
unjudged_count,
|
||||
)
|
||||
return BookJudgmentResult()
|
||||
|
||||
mentions = index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0
|
||||
session.commit()
|
||||
except Exception:
|
||||
session.rollback()
|
||||
logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id)
|
||||
return BookJudgmentResult(failed=True)
|
||||
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s "
|
||||
"duration_ms=%.1f",
|
||||
source_id,
|
||||
judged,
|
||||
len(protected),
|
||||
mentions,
|
||||
(perf_counter() - book_started_at) * 1000,
|
||||
)
|
||||
return BookJudgmentResult(judged=judged, protected=len(protected), mentions=mentions, committed=True)
|
||||
|
||||
|
||||
def judge_candidate_phrases_for_book(
|
||||
session: Session,
|
||||
book_id: int,
|
||||
series_id: int | None,
|
||||
normalized_book_text: str,
|
||||
config: EbookSearchConfig,
|
||||
) -> tuple[int, list[EbookProtectedPhrase]]:
|
||||
"""Judge unjudged candidate phrase rows for one book.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book whose candidates are judged.
|
||||
series_id (int | None): Series scope for promoted protected phrases.
|
||||
normalized_book_text (str): Whole book text, already normalized, for context lookups.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
tuple[int, list[EbookProtectedPhrase]]: Number of candidates judged and the promoted phrases.
|
||||
"""
|
||||
judgment_limit = config.protected_phrase_llm_candidates_per_book
|
||||
if judgment_limit <= 0:
|
||||
logger.info("ebook_candidate_phrase_judgment_skipped_llm_limit_zero book_id=%s", book_id)
|
||||
return 0, []
|
||||
|
||||
existing_protected = count_protected_phrases(session, book_id)
|
||||
target_remaining: int | None = None
|
||||
if config.phrase_target_protected_per_book > 0:
|
||||
target_remaining = max(config.phrase_target_protected_per_book - existing_protected, 0)
|
||||
if target_remaining == 0:
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_skipped_target_met book_id=%s existing_protected=%s target=%s",
|
||||
book_id,
|
||||
existing_protected,
|
||||
config.phrase_target_protected_per_book,
|
||||
)
|
||||
return 0, []
|
||||
|
||||
rows = load_candidates_for_judgment(session, book_id, judgment_limit, config)
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_candidates_loaded book_id=%s candidates=%s existing_protected=%s "
|
||||
"target_remaining=%s judgment_limit=%s",
|
||||
book_id,
|
||||
len(rows),
|
||||
existing_protected,
|
||||
target_remaining,
|
||||
judgment_limit,
|
||||
)
|
||||
|
||||
judged_count = 0
|
||||
protected: list[EbookProtectedPhrase] = []
|
||||
for row_number, row in enumerate(rows, start=1):
|
||||
candidate, judgment, candidate_row = judge_candidate_row(
|
||||
session, book_id, series_id, normalized_book_text, row, row_number, len(rows), config
|
||||
)
|
||||
judged_count += 1
|
||||
if not should_protect_judged_candidate(row, candidate, judgment, book_id, config):
|
||||
continue
|
||||
protected.append(upsert_protected_phrase(session, book_id, series_id, candidate, judgment, candidate_row))
|
||||
if target_remaining is not None and len(protected) >= target_remaining:
|
||||
break
|
||||
|
||||
session.flush()
|
||||
return judged_count, protected
|
||||
|
||||
|
||||
def judge_candidate_row(
|
||||
session: Session,
|
||||
book_id: int,
|
||||
series_id: int | None,
|
||||
normalized_book_text: str,
|
||||
row: EbookCandidatePhrase,
|
||||
row_number: int,
|
||||
total_rows: int,
|
||||
config: EbookSearchConfig,
|
||||
) -> tuple[PhraseCandidate, LLMJudgment, EbookCandidatePhrase]:
|
||||
"""Run and persist the LLM judgment for a single candidate row.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book the candidate belongs to.
|
||||
series_id (int | None): Series scope for the saved candidate.
|
||||
normalized_book_text (str): Whole book text, already normalized, for context lookups.
|
||||
row (EbookCandidatePhrase): Stored candidate row to judge.
|
||||
row_number (int): 1-based position of the row in the batch, for logging.
|
||||
total_rows (int): Total rows in the batch, for logging.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
tuple[PhraseCandidate, LLMJudgment, EbookCandidatePhrase]: The candidate, its judgment,
|
||||
and the persisted candidate row.
|
||||
"""
|
||||
row_started_at = perf_counter()
|
||||
candidate = phrase_candidate_from_row(row)
|
||||
candidate.sample_contexts = row.sample_contexts or get_sample_contexts(normalized_book_text, candidate.phrase_norm)
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_candidate_start book_id=%s candidate_id=%s row_number=%s rows=%s "
|
||||
"phrase=%r score=%.3f raw_count=%s chapter_count=%s",
|
||||
book_id,
|
||||
row.id,
|
||||
row_number,
|
||||
total_rows,
|
||||
candidate.phrase_norm,
|
||||
candidate.candidate_score,
|
||||
candidate.raw_count,
|
||||
candidate.chapter_count,
|
||||
)
|
||||
judgment = judge_candidate_with_llm(candidate, config)
|
||||
candidate_row = save_candidate_to_db(session, book_id, series_id, candidate, judgment=judgment)
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_candidate_complete book_id=%s candidate_id=%s phrase=%r keep=%s "
|
||||
"confidence=%.3f importance=%.3f category=%r duration_ms=%.1f",
|
||||
book_id,
|
||||
row.id,
|
||||
candidate.phrase_norm,
|
||||
judgment.keep,
|
||||
judgment.confidence,
|
||||
judgment.importance,
|
||||
judgment.category,
|
||||
(perf_counter() - row_started_at) * 1000,
|
||||
)
|
||||
return candidate, judgment, candidate_row
|
||||
|
||||
|
||||
def should_protect_judged_candidate(
|
||||
row: EbookCandidatePhrase,
|
||||
candidate: PhraseCandidate,
|
||||
judgment: LLMJudgment,
|
||||
book_id: int,
|
||||
config: EbookSearchConfig,
|
||||
) -> bool:
|
||||
"""Report whether a judged candidate qualifies to become a protected phrase.
|
||||
|
||||
Args:
|
||||
row (EbookCandidatePhrase): Stored candidate row the judgment came from.
|
||||
candidate (PhraseCandidate): In-memory candidate that was judged.
|
||||
judgment (LLMJudgment): Judge decision for the candidate.
|
||||
book_id (int): Book the candidate belongs to, for logging.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings.
|
||||
|
||||
Returns:
|
||||
bool: True when the judged candidate should be promoted to a protected phrase.
|
||||
"""
|
||||
if not judgment.keep or judgment.confidence < config.protected_phrase_confidence_threshold:
|
||||
return False
|
||||
accepted_norm = normalize_text(judgment.canonical or candidate.phrase_text)
|
||||
accepted_token_count = len(accepted_norm.split())
|
||||
if accepted_token_count < config.phrase_min_tokens:
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_candidate_skip_short_canonical book_id=%s candidate_id=%s "
|
||||
"phrase=%r canonical=%r token_count=%s min_tokens=%s",
|
||||
book_id,
|
||||
row.id,
|
||||
candidate.phrase_norm,
|
||||
accepted_norm,
|
||||
accepted_token_count,
|
||||
config.phrase_min_tokens,
|
||||
)
|
||||
return False
|
||||
if is_most_common_word_phrase(accepted_norm):
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_judgment_candidate_skip_common_canonical book_id=%s candidate_id=%s "
|
||||
"phrase=%r canonical=%r",
|
||||
book_id,
|
||||
row.id,
|
||||
candidate.phrase_norm,
|
||||
accepted_norm,
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
"""LLM judging of extracted candidate phrases."""
|
||||
|
||||
|
||||
def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConfig) -> LLMJudgment:
|
||||
"""Ask the configured chat model to judge one pre-extracted candidate.
|
||||
|
||||
Args:
|
||||
candidate (PhraseCandidate): Candidate to send to the LLM judge.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings and chat configuration.
|
||||
|
||||
Returns:
|
||||
LLMJudgment: The parsed structured judgment for the candidate.
|
||||
"""
|
||||
payload = {
|
||||
"phrase": candidate.phrase_norm,
|
||||
"token_count": candidate.token_count,
|
||||
"sources": candidate_source_names(candidate),
|
||||
"raw_count": candidate.raw_count,
|
||||
"chapter_count": candidate.chapter_count,
|
||||
"contexts": candidate.sample_contexts,
|
||||
}
|
||||
messages = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"Judge whether a candidate phrase from a book should be protected for RAG retrieval. "
|
||||
"Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific phrases, "
|
||||
"unstable fragments, and phrases kept only because they are frequent. Keep people, places, "
|
||||
"organizations, factions, events, technologies, fictional conditions, magic systems, formal titles, "
|
||||
"named concepts, and recurring world-specific terms. Return only a JSON object with keys: keep, "
|
||||
"canonical, category, aliases, confidence, importance, allow_nested, suppress_children, reason."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": json.dumps(payload, ensure_ascii=True)},
|
||||
]
|
||||
return parse_llm_judgment(request_chat_completion(config, messages), config)
|
||||
|
||||
|
||||
def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment:
|
||||
"""Parse and validate an LLM phrase-judge response.
|
||||
|
||||
Args:
|
||||
content (str): Raw model response text.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings supplying nesting defaults.
|
||||
|
||||
Returns:
|
||||
LLMJudgment: The parsed and validated judgment.
|
||||
|
||||
Raises:
|
||||
TypeError: If the decoded JSON body is not an object.
|
||||
"""
|
||||
body = json.loads(extract_json_object(content))
|
||||
if not isinstance(body, dict):
|
||||
msg = "LLM phrase judge response is not a JSON object"
|
||||
raise TypeError(msg)
|
||||
|
||||
aliases = body.get("aliases", ())
|
||||
if not isinstance(aliases, list | tuple):
|
||||
aliases = ()
|
||||
return LLMJudgment(
|
||||
keep=bool(body.get("keep", False)),
|
||||
canonical=optional_text(body.get("canonical")),
|
||||
category=optional_text(body.get("category")),
|
||||
aliases=tuple(str(alias) for alias in aliases if isinstance(alias, str) and alias.strip()),
|
||||
confidence=clamped_float(body.get("confidence"), default=0.0),
|
||||
importance=clamped_float(body.get("importance"), default=0.5),
|
||||
allow_nested=bool(body.get("allow_nested", config.phrase_default_allow_nested)),
|
||||
suppress_children=bool(body.get("suppress_children", config.phrase_default_suppress_children)),
|
||||
reason=optional_text(body.get("reason")),
|
||||
)
|
||||
|
||||
|
||||
def extract_json_object(content: str) -> str:
|
||||
"""Extract a JSON object from plain or fenced model output.
|
||||
|
||||
Args:
|
||||
content (str): Raw model response text.
|
||||
|
||||
Returns:
|
||||
str: The substring spanning the first JSON object.
|
||||
|
||||
Raises:
|
||||
ValueError: If no JSON object is found in the response.
|
||||
"""
|
||||
stripped = content.strip()
|
||||
if stripped.startswith("{") and stripped.endswith("}"):
|
||||
return stripped
|
||||
match = JSON_OBJECT_RE.search(stripped)
|
||||
if match is None:
|
||||
msg = "LLM phrase judge response did not contain a JSON object"
|
||||
raise ValueError(msg)
|
||||
return match.group(0)
|
||||
|
||||
|
||||
def optional_text(value: object) -> str | None:
|
||||
"""Return stripped text for a nullable JSON value.
|
||||
|
||||
Args:
|
||||
value (object): Decoded JSON value that may or may not be a string.
|
||||
|
||||
Returns:
|
||||
str | None: The stripped string, or ``None`` when it is not a non-empty string.
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
stripped = value.strip()
|
||||
return stripped or None
|
||||
|
||||
|
||||
def clamped_float(value: object, *, default: float) -> float:
|
||||
"""Coerce a JSON number into the 0.0 to 1.0 range.
|
||||
|
||||
Args:
|
||||
value (object): Decoded JSON value that may or may not be a number.
|
||||
default (float): Fallback returned when ``value`` is not numeric.
|
||||
|
||||
Returns:
|
||||
float: The value clamped to ``[0.0, 1.0]``, or ``default`` when non-numeric.
|
||||
"""
|
||||
if not isinstance(value, int | float):
|
||||
return default
|
||||
return min(max(float(value), 0.0), 1.0)
|
||||
Reference in New Issue
Block a user