"""Book-level orchestration for LLM judging and promotion of candidate phrases.""" from __future__ import annotations import asyncio import json import logging import re from time import perf_counter from typing import TYPE_CHECKING import httpx from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession 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_junk_phrase, is_most_common_word_phrase, score_candidate, ) 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 collections.abc import Sequence from sqlalchemy.ext.asyncio import AsyncEngine from python.ebook_search.config import EbookSearchConfig from python.ebook_search.protected_phrases.models import PhraseCandidate from python.orm.richie import EbookProtectedPhrase JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL) logger = logging.getLogger(__name__) async def judge_candidate_phrases_for_books( engine: AsyncEngine, config: EbookSearchConfig, *, source_ids: Sequence[int] | None = None, ) -> PhraseJudgmentBackfillResult: """Judge candidate phrases for books, fanning LLM calls out across books and phrases. Up to ``phrase_judge_book_workers`` books are judged at once, and within each book candidates are judged in concurrent chunks of ``phrase_judge_phrase_workers``. Each book uses its own short-lived sessions for reads and writes; no database connection is held while LLM calls are in flight. For a pseudo-single-threaded run (solo testing, debugging), set both worker settings to 1. Args: engine (AsyncEngine): Engine used to open one session per book. config (EbookSearchConfig): Runtime phrase-tuning settings and chat configuration. source_ids (Sequence[int] | None): Books to judge; ``None`` judges every indexed book. Returns: PhraseJudgmentBackfillResult: Per-corpus counts of books judged, failures, candidates, protected phrases, and mentions. """ if source_ids is None: async with AsyncSession(engine) as session: source_ids = list((await session.scalars(select(EbookSource.id).order_by(EbookSource.id))).all()) books_seen = len(source_ids) book_workers = max(1, config.phrase_judge_book_workers) phrase_workers = max(1, config.phrase_judge_phrase_workers) logger.info( "ebook_candidate_phrase_judgment_start books_seen=%s book_workers=%s phrase_workers=%s " "confidence_threshold=%.2f", books_seen, book_workers, phrase_workers, config.protected_phrase_confidence_threshold, ) book_semaphore = asyncio.Semaphore(book_workers) max_connections = book_workers * phrase_workers limits = httpx.Limits(max_connections=max_connections, max_keepalive_connections=max_connections) async with httpx.AsyncClient(limits=limits) as client: outcomes = await asyncio.gather( *(judge_one_book_async(engine, source_id, config, client, book_semaphore) 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 async def judge_one_book_async( engine: AsyncEngine, source_id: int, config: EbookSearchConfig, client: httpx.AsyncClient, book_semaphore: asyncio.Semaphore, ) -> BookJudgmentResult: """Judge one book concurrently and persist the outcome, honoring the book-level limit. Args: engine (AsyncEngine): Engine used to open the book's read and write sessions. source_id (int): Book to judge candidates for. config (EbookSearchConfig): Runtime phrase-tuning settings. client (httpx.AsyncClient): Shared async client for LLM calls. book_semaphore (asyncio.Semaphore): Caps how many books judge at once. Returns: BookJudgmentResult: The book's judgment outcome. """ async with book_semaphore: try: prepared = await prepare_book_judgment(engine, source_id, config) if prepared is None: return BookJudgmentResult() work_items, target_remaining = prepared judged = await judge_book_candidates_async(client, config, source_id, work_items, target_remaining) if not judged: return BookJudgmentResult() return await persist_book_judgments(engine, source_id, config, judged) except Exception: logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id) return BookJudgmentResult(failed=True) async def prepare_book_judgment( engine: AsyncEngine, source_id: int, config: EbookSearchConfig, ) -> tuple[list[tuple[int, PhraseCandidate]], int | None] | None: """Load one book's candidates to judge, with sample contexts, on a short-lived read session. Args: engine (AsyncEngine): Engine used to open the read session. source_id (int): Book to load candidates for. config (EbookSearchConfig): Runtime phrase-tuning settings. Returns: tuple[list[tuple[int, PhraseCandidate]], int | None] | None: Candidate rows paired with in-memory candidates and the remaining protected-phrase target, or ``None`` when the book has nothing to judge. """ judgment_limit = config.protected_phrase_llm_candidates_per_book if judgment_limit <= 0: return None async with AsyncSession(engine) as session: if not await count_unjudged_candidates(session, source_id, config): logger.info("ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s", source_id) return None existing_protected = await count_protected_phrases(session, source_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 source_id=%s existing_protected=%s target=%s", source_id, existing_protected, config.phrase_target_protected_per_book, ) return None book_text = await load_book_text(session, source_id) if not book_text: logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id) return None normalized_book_text = normalize_text(book_text) # Stored rows may predate the current junk filters and score weights, so re-filter and # rescore every unjudged row here instead of trusting the persisted candidate_score. rows = await load_candidates_for_judgment(session, source_id, config) scored_items: list[tuple[int, PhraseCandidate]] = [] skipped_junk = 0 for row in rows: candidate = phrase_candidate_from_row(row) if is_junk_phrase(candidate.phrase_norm.split()): skipped_junk += 1 continue candidate.candidate_score = score_candidate(candidate, config) scored_items.append((row.id, candidate)) scored_items.sort(key=lambda item: item[1].candidate_score, reverse=True) work_items = scored_items[:judgment_limit] for _, candidate in work_items: candidate.sample_contexts = candidate.sample_contexts or get_sample_contexts( normalized_book_text, candidate.phrase_norm ) logger.info( "ebook_candidate_phrase_judgment_candidates_loaded source_id=%s candidates=%s skipped_junk=%s " "unjudged_rows=%s existing_protected=%s target_remaining=%s judgment_limit=%s", source_id, len(work_items), skipped_junk, len(rows), existing_protected, target_remaining, judgment_limit, ) return work_items, target_remaining async def judge_book_candidates_async( client: httpx.AsyncClient, config: EbookSearchConfig, source_id: int, work_items: list[tuple[int, PhraseCandidate]], target_remaining: int | None, ) -> list[tuple[int, PhraseCandidate, LLMJudgment, bool]]: """Judge a book's candidates in concurrent chunks, stopping once the target is reached. Promotion decisions are made in memory so judging can stop early without any database writes. Args: client (httpx.AsyncClient): Shared async client for LLM calls. config (EbookSearchConfig): Runtime phrase-tuning settings. source_id (int): Book being judged, for logging. work_items (list[tuple[int, PhraseCandidate]]): Candidate row ids paired with candidates, in best-first score order. target_remaining (int | None): Remaining protected-phrase target, or ``None`` for no cap. Returns: list[tuple[int, PhraseCandidate, LLMJudgment, bool]]: Judged rows with their judgment and whether each should be promoted. """ chunk_size = max(1, config.phrase_judge_phrase_workers) judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]] = [] promoted = 0 for start in range(0, len(work_items), chunk_size): chunk = work_items[start : start + chunk_size] judgments = await asyncio.gather(*(judge_candidate_async(client, config, candidate) for _, candidate in chunk)) for (candidate_id, candidate), judgment in zip(chunk, judgments, strict=True): promote = (target_remaining is None or promoted < target_remaining) and should_protect_judged_candidate( candidate, judgment, source_id, config, candidate_id=candidate_id ) if promote: promoted += 1 judged.append((candidate_id, candidate, judgment, promote)) if target_remaining is not None and promoted >= target_remaining: break return judged async def judge_candidate_async( client: httpx.AsyncClient, config: EbookSearchConfig, candidate: PhraseCandidate, ) -> LLMJudgment: """Judge one candidate with the LLM over the shared async client. Args: client (httpx.AsyncClient): Shared async client for LLM calls. config (EbookSearchConfig): Runtime phrase-tuning settings. candidate (PhraseCandidate): Candidate to judge. Returns: LLMJudgment: The parsed judgment. """ content = await request_chat_completion(client, config, build_judge_messages(candidate)) return parse_llm_judgment(content, config) async def persist_book_judgments( engine: AsyncEngine, source_id: int, config: EbookSearchConfig, judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]], ) -> BookJudgmentResult: """Persist one book's judgments and promotions in a single committed transaction. Args: engine (AsyncEngine): Engine used to open the write session. source_id (int): Book being persisted. config (EbookSearchConfig): Runtime phrase-tuning settings. judged (list[tuple[int, PhraseCandidate, LLMJudgment, bool]]): Judged candidates with their judgment and promotion flag. Returns: BookJudgmentResult: The book's committed counts, or a failed result on error. """ book_started_at = perf_counter() async with AsyncSession(engine, expire_on_commit=False) as session: try: protected: list[EbookProtectedPhrase] = [] for candidate_id, candidate, judgment, promote in judged: candidate_row = await save_candidate_to_db(session, source_id, None, candidate, judgment=judgment) if promote: protected.append( await upsert_protected_phrase(session, source_id, None, candidate, judgment, candidate_row) ) logger.info( "ebook_candidate_phrase_judgment_candidate_complete source_id=%s candidate_id=%s phrase=%r " "keep=%s confidence=%.3f category=%r promoted=%s", source_id, candidate_id, candidate.phrase_norm, judgment.keep, judgment.confidence, judgment.category, promote, ) await session.flush() mentions = await index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0 await session.commit() except Exception: await session.rollback() logger.exception("ebook_candidate_phrase_judgment_book_persist_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, len(judged), len(protected), mentions, (perf_counter() - book_started_at) * 1000, ) return BookJudgmentResult(judged=len(judged), protected=len(protected), mentions=mentions, committed=True) def should_protect_judged_candidate( candidate: PhraseCandidate, judgment: LLMJudgment, book_id: int, config: EbookSearchConfig, *, candidate_id: int, ) -> bool: """Report whether a judged candidate qualifies to become a protected phrase. Args: 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. candidate_id (int): Stored candidate row id the judgment came from, for logging. 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_tokens = accepted_norm.split() accepted_token_count = len(accepted_tokens) 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, candidate_id, candidate.phrase_norm, accepted_norm, accepted_token_count, config.phrase_min_tokens, ) return False if is_most_common_word_phrase(accepted_tokens): logger.info( "ebook_candidate_phrase_judgment_candidate_skip_common_canonical book_id=%s candidate_id=%s " "phrase=%r canonical=%r", book_id, candidate_id, candidate.phrase_norm, accepted_norm, ) return False return True def build_judge_messages(candidate: PhraseCandidate) -> list[dict[str, str]]: """Build the chat messages used to judge one candidate phrase. Args: candidate (PhraseCandidate): Candidate to describe for the judge. Returns: list[dict[str, str]]: OpenAI-style system and user messages. """ 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, } return [ { "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)}, ] 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)