dotfiles/python/ebook_search/protected_phrases/judge_ngrams.py

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