"""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)