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 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(
        f"ebook_candidate_phrase_judgment_start {books_seen=} {book_workers=} {phrase_workers=} "
        f"{config.protected_phrase_confidence_threshold=:.2f}"
    )

    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(
        f"ebook_candidate_phrase_judgment_complete {result.books_seen=} {result.books_judged=} {result.books_failed=} "
        f"{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(f"ebook_candidate_phrase_judgment_book_failed {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(f"ebook_candidate_phrase_judgment_book_skip_no_unjudged {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(
                    f"ebook_candidate_phrase_judgment_skipped_target_met {source_id=} {existing_protected=} "
                    f"{config.phrase_target_protected_per_book=}"
                )
                return None
        book_text = await load_book_text(session, source_id)
        if not book_text:
            logger.warning(f"ebook_candidate_phrase_judgment_book_empty {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(
        f"ebook_candidate_phrase_judgment_candidates_loaded {source_id=} candidates={len(work_items)} {skipped_junk=} "
        f"unjudged_rows={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(
                    f"ebook_candidate_phrase_judgment_candidate_complete {source_id=} {candidate_id=} "
                    f"{candidate.phrase_norm=} {judgment.keep=} {judgment.confidence=:.3f} {judgment.category=} "
                    f"{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(f"ebook_candidate_phrase_judgment_book_persist_failed {source_id=}")
            return BookJudgmentResult(failed=True)
    logger.info(
        f"ebook_candidate_phrase_judgment_book_committed {source_id=} judged={len(judged)} protected={len(protected)} "
        f"{mentions=} duration_ms={(perf_counter() - book_started_at) * 1000:.1f}"
    )
    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(
            f"ebook_candidate_phrase_judgment_candidate_skip_short_canonical {book_id=} {candidate_id=} "
            f"{candidate.phrase_norm=} {accepted_norm=} {accepted_token_count=} {config.phrase_min_tokens=}"
        )
        return False
    if is_most_common_word_phrase(accepted_tokens):
        logger.info(
            f"ebook_candidate_phrase_judgment_candidate_skip_common_canonical {book_id=} {candidate_id=} "
            f"{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)