dotfiles/python/ebook_search/protected_phrases/models.py

"""Dataclasses shared by protected phrase extraction, judging, matching, and backfills."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Mapping

    from python.orm.richie import EbookProtectedPhrase


@dataclass(slots=True)
class PhraseCandidate:
    """A phrase candidate with merged extraction-source metadata.

    Attributes:
        phrase_text (str): Display text for the phrase.
        phrase_norm (str): Normalized phrase used as the merge key.
        token_count (int): Number of normalized tokens in the phrase.
        source_raw_ngram (bool): Whether the raw n-gram extractor produced the phrase.
        source_yake (bool): Whether YAKE keyword extraction produced the phrase.
        source_capitalized (bool): Whether the capitalized-run extractor produced the phrase.
        source_metadata (bool): Whether book metadata produced the phrase.
        raw_count (int): Occurrences counted across the book text.
        chapter_count (int): Number of chapters containing the phrase.
        yake_score (float | None): Raw YAKE score when available; lower is better.
        candidate_score (float): Combined pre-judging score.
        sample_contexts (list[str]): Normalized context snippets around occurrences.
    """

    phrase_text: str
    phrase_norm: str
    token_count: int
    source_raw_ngram: bool = False
    source_yake: bool = False
    source_capitalized: bool = False
    source_metadata: bool = False
    raw_count: int = 0
    chapter_count: int = 0
    yake_score: float | None = None
    candidate_score: float = 0.0
    sample_contexts: list[str] = field(default_factory=list)


@dataclass(frozen=True, slots=True)
class LLMJudgment:
    """A structured phrase judgment returned by the LLM judge.

    Attributes:
        keep (bool): Whether the judge accepted the phrase for protection.
        canonical (str | None): Canonical phrase text chosen by the judge.
        category (str | None): Phrase category such as person, place, or event.
        aliases (tuple[str, ...]): Alternate surface forms for the phrase.
        confidence (float): Judge confidence between 0.0 and 1.0.
        importance (float): Judge importance between 0.0 and 1.0.
        allow_nested (bool): Whether the phrase may match inside a larger kept match.
        suppress_children (bool): Whether the phrase suppresses matches nested inside it.
        reason (str | None): Free-text explanation from the judge.
    """

    keep: bool
    canonical: str | None
    category: str | None
    aliases: tuple[str, ...]
    confidence: float
    importance: float = 0.5
    allow_nested: bool = False
    suppress_children: bool = True
    reason: str | None = None


@dataclass(frozen=True, slots=True)
class PhraseLookup:
    """In-memory phrase metadata used for constant-time text-window checks.

    Attributes:
        phrase_ids_by_norm (Mapping[str, tuple[int, ...]]): Canonical and alias norms to phrase ids.
        phrases_by_id (Mapping[int, EbookProtectedPhrase]): Protected phrase metadata by id.
        min_tokens (int): Smallest token-window size to test.
        max_tokens (int): Largest token-window size to test.
    """

    phrase_ids_by_norm: Mapping[str, tuple[int, ...]]
    phrases_by_id: Mapping[int, EbookProtectedPhrase]
    min_tokens: int
    max_tokens: int


@dataclass(frozen=True, slots=True)
class PhraseMatch:
    """A detected phrase match with protected-phrase metadata attached.

    Attributes:
        phrase_id (int): Protected phrase id.
        matched_norm (str): Normalized window text that matched.
        phrase_text (str): Display text of the protected phrase.
        phrase_norm (str): Normalized text of the protected phrase.
        canonical_id (str): Deterministic ``category:slug`` identifier.
        phrase_type (str | None): Phrase category.
        token_count (int): Number of tokens in the match.
        confidence (float): Stored judge confidence.
        importance (float): Stored judge importance.
        allow_nested (bool): Whether the phrase may match inside a larger kept match.
        suppress_children (bool): Whether the phrase suppresses matches nested inside it.
        start_token (int): Index of the first matched token.
        end_token (int): Index one past the last matched token.
        start_char (int | None): Start character offset in the source text.
        end_char (int | None): End character offset in the source text.
        book_id (int | None): Book scope of the phrase.
        series_id (int | None): Series scope of the phrase.
    """

    phrase_id: int
    matched_norm: str
    phrase_text: str
    phrase_norm: str
    canonical_id: str
    phrase_type: str | None
    token_count: int
    confidence: float
    importance: float
    allow_nested: bool
    suppress_children: bool
    start_token: int
    end_token: int
    start_char: int | None = None
    end_char: int | None = None
    book_id: int | None = None
    series_id: int | None = None


@dataclass(frozen=True, slots=True)
class PhraseCandidateGenerationResult:
    """Summary of candidate phrase extraction for indexed books.

    Attributes:
        books_seen (int): Indexed books examined.
        books_built (int): Books that had candidates generated and committed.
        candidate_phrases (int): Candidate phrases stored across all books.
    """

    books_seen: int
    books_built: int
    candidate_phrases: int


@dataclass(frozen=True, slots=True)
class CorpusPhraseStats:
    """Corpus-wide candidate and protected phrase counts for the admin page.

    Attributes:
        total_books (int): Indexed books in the corpus.
        books_with_candidates (int): Books that have candidate phrases generated.
        books_fully_judged (int): Books with candidates where every candidate has been judged.
        candidate_phrases (int): Candidate phrases stored across all books.
        judged_candidates (int): Candidate phrases that have been LLM judged.
        unjudged_candidates (int): Candidate phrases still waiting for judgment.
        protected_phrases (int): Protected phrases promoted across all books.
    """

    total_books: int
    books_with_candidates: int
    books_fully_judged: int
    candidate_phrases: int
    judged_candidates: int
    unjudged_candidates: int
    protected_phrases: int


@dataclass(frozen=True, slots=True)
class PhraseJudgmentBackfillResult:
    """Summary of LLM judging for stored candidate phrases.

    Attributes:
        books_seen (int): Indexed books examined.
        books_judged (int): Books with judgments committed.
        books_failed (int): Books rolled back after an error.
        candidates_judged (int): Candidate phrases sent to the LLM judge.
        protected_phrases (int): Protected phrases promoted from candidates.
        phrase_mentions (int): Chunk phrase mentions indexed across all books.
    """

    books_seen: int
    books_judged: int
    books_failed: int
    candidates_judged: int
    protected_phrases: int
    phrase_mentions: int


@dataclass(frozen=True, slots=True)
class BookJudgmentResult:
    """Outcome of judging one book's candidate phrases.

    Attributes:
        judged (int): Candidate phrases sent to the LLM judge.
        protected (int): Protected phrases promoted from candidates.
        mentions (int): Chunk phrase mentions indexed for the book.
        committed (bool): Whether the book's judgments were committed.
        failed (bool): Whether the book was rolled back after an error.
    """

    judged: int = 0
    protected: int = 0
    mentions: int = 0
    committed: bool = False
    failed: bool = False


@dataclass(frozen=True, slots=True)
class BookCandidateResult:
    """Outcome of generating one book's candidate phrases.

    Attributes:
        candidates (int): Candidate phrases stored for the book.
        built (bool): Whether candidate generation was committed.
    """

    candidates: int = 0
    built: bool = False


@dataclass(frozen=True, slots=True)
class PhraseRecalculationResult:
    """Summary of phrase cleanup and candidate regeneration for one book.

    Attributes:
        book_id (int): Book the recalculation ran against.
        deleted_candidates (int): Candidate phrase rows deleted.
        deleted_protected_phrases (int): Protected phrase rows deleted.
        deleted_aliases (int): Phrase alias rows deleted.
        deleted_mentions (int): Chunk phrase mention rows deleted.
        candidate_phrases (int): Candidate phrases regenerated after cleanup.
    """

    book_id: int
    deleted_candidates: int
    deleted_protected_phrases: int
    deleted_aliases: int
    deleted_mentions: int
    candidate_phrases: int