arborist/aborist/qa/frame.py

"""Reference-frame detection for claim-lattice queries.

Closes the answer-side gap left open by the phrase-pattern retrieval
route (commit `1b8677d`). The
phrase route surfaces the right article for allusion-shape queries
(e.g. Nineteen Eighty-Four for "has oceania always been at war with
east asia"); this module then classifies the answer's intended frame
so the prompt + renderer can produce a multi-frame answer that
distinguishes literal from reference from in-universe-propaganda.

See `docs/tickets/ticket-000002-reference-frame-polarity-contract.md` for the
full scope. This module is the *substrate* — it returns a
`FrameDetection` per call. Prompt augmentation + renderer extension
live in their respective layers and consume this output.

Hard or soft? **Soft.** This is a sidecar-class signal that informs
prompt augmentation + renderer presentation; it never enters the
proof path or `cache_key`. Per CLAUDE.md "Soft hash vs hard hash":
heuristics, scores, and judgment calls live in the soft channel.
The hard verifier never consults `FrameDetection` to accept or
reject claims.
"""
from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Literal


# Parenthetical disambiguation suffixes Wikipedia uses for fictional /
# reference works. Case-insensitive; matches at the open-paren without
# a leading word-boundary (`\b` doesn't match between space and `(`
# since both are non-word chars).
_REFERENCE_TITLE_SUFFIX_RE = re.compile(
    r"\("
    r"(?:novel|novella|short\s+story|story|film|movie|play|"
    r"video\s+game|game|comic|comics|graphic\s+novel|book|"
    r"franchise|series|tv\s+series|television\s+series|"
    r"manga|anime|fictional\s+character|character|episode|"
    r"album|song|musical|opera|poem)"
    r"\)",
    re.IGNORECASE,
)


# Body-side fiction markers. When a retrieved source's body sample
# contains a high density of these tokens, the source is a reference
# work even when the title lacks a disambiguation suffix
# (Nineteen_Eighty-Four article being the canonical case — its title
# is the bare work-name with no parenthetical, but the body opens
# with "is a dystopian social science fiction novel by ...").
_FICTION_BODY_MARKERS = frozenset({
    "novel", "novels", "novella", "novellas", "fiction", "fictional",
    "fantasy", "dystopian", "satire", "satirical", "allegory",
    "allegorical", "characters", "protagonist", "antagonist",
    "plot", "narrator", "narrative", "chapter", "chapters",
    "published", "publishing", "publisher", "screenplay", "directed",
    "starring", "cast", "filmed", "broadcast", "aired", "premiered",
    "written", "writer", "author", "playwright", "screenwriter",
    "adaptation", "adapted", "sequel", "prequel", "trilogy",
    "saga", "franchise", "manga", "anime",
})

# Distinct-marker threshold: a body sample must contain at least
# this many DISTINCT fiction markers to classify as reference work.
# Counting distinct markers (not total occurrences) defends against
# a single common-vocabulary use repeated through the sample (e.g.
# "novel approach" + "novel approach" doesn't count as two markers,
# only one distinct).
_FICTION_DISTINCT_MARKER_THRESHOLD = 3


FrameKind = Literal["literal", "reference", "ambiguous", "no_phrase_route"]


@dataclass(frozen=True)
class FrameDetection:
    """Reference-frame classification for a claim-lattice query.

    Returned by ``detect_frame``. Consumed by:
      - prompt augmentation (when ``frame_kind == "reference"``,
        inject a polarity instruction into the system prompt)
      - renderer (when ``frame_kind == "reference"``, render a
        Frame notes section under the answer)

    Never persisted in `providence_cache`, never enters
    `governance_policy_hash`, never folds into the run-DAG. Soft
    sidecar signal only.
    """
    frame_kind: FrameKind
    reference_title: str | None = None      # e.g. "Nineteen Eighty-Four"
    reference_uri: str | None = None        # e.g. wiki URI tail
    confidence: float = 0.0                 # 0.0-1.0, heuristic-based


def _title_indicates_reference_work(title: str | None) -> bool:
    """Title carries a Wikipedia-style disambiguation suffix for a
    fictional / reference work."""
    if not title:
        return False
    return bool(_REFERENCE_TITLE_SUFFIX_RE.search(title))


def _body_indicates_reference_work(body_sample: str | None) -> bool:
    """Body sample contains at least
    ``_FICTION_DISTINCT_MARKER_THRESHOLD`` distinct fiction markers.
    `body_sample` should be a short prefix (~1-2 KB) of the source's
    first chunk; we don't scan the whole article.

    Counting DISTINCT markers (not total occurrences) keeps the
    heuristic robust against single-marker repetition: a history
    article saying "novel approach" twice doesn't trip the detector
    because only one marker (`novel`) fired. A reference-work lead
    typically clusters several distinct markers (`novel` + `published`
    + `characters` + `plot` + `protagonist` etc.).
    """
    if not body_sample:
        return False
    text_lower = body_sample.lower()
    distinct_hits = 0
    for marker in _FICTION_BODY_MARKERS:
        pattern = r"\b" + re.escape(marker) + r"\b"
        if re.search(pattern, text_lower):
            distinct_hits += 1
            if distinct_hits >= _FICTION_DISTINCT_MARKER_THRESHOLD:
                return True
    return False


def _source_is_reference_work(source: dict) -> bool:
    """A source is a reference work when its title carries a
    disambiguation suffix OR its body opens with high fiction-marker
    density. Either signal alone suffices.
    """
    if _title_indicates_reference_work(source.get("title")):
        return True
    return _body_indicates_reference_work(source.get("body_sample"))


def detect_frame(
    question: str,
    sources: list[dict],
    phrase_match_roots: set[str] | None = None,
) -> FrameDetection:
    """Classify the query's intended frame.

    Inputs:
      - ``question`` — the user's question text (currently unused
        directly; kept for future shape-detection extensions).
      - ``sources`` — the retrieved sources for the query. Each
        dict should carry ``document_root``, ``title``, and
        ``body_sample`` (a short prefix of the first chunk).
      - ``phrase_match_roots`` — set of document_roots that were
        surfaced by the phrase-pattern retrieval route (allusion
        / verbatim multi-token sequence match). When non-empty,
        the query is allusion-shape and the corresponding sources
        are candidate reference works.

    Algorithm:
      1. If no phrase route fired (``phrase_match_roots`` empty),
         the query is literal — no allusion signal at retrieval.
      2. If phrase route fired AND any phrase-matched source is a
         reference work (title suffix or body fiction-density), the
         query is reference-frame; pick the first such source as
         the named reference.
      3. If phrase route fired but no matched source is a reference
         work, classify ``ambiguous`` — the verbatim phrase appeared
         somewhere in the corpus but not in a fictional context.

    Returns:
        FrameDetection with the classified ``frame_kind``,
        ``reference_title`` / ``reference_uri`` populated for
        reference-frame queries, and a confidence score that
        consumers can threshold if they want a more conservative
        gate.
    """
    if not phrase_match_roots:
        return FrameDetection(frame_kind="no_phrase_route", confidence=0.0)
    candidate: dict | None = None
    for src in sources:
        if src.get("document_root") not in phrase_match_roots:
            continue
        if _source_is_reference_work(src):
            candidate = src
            break
    if candidate is not None:
        return FrameDetection(
            frame_kind="reference",
            reference_title=candidate.get("title"),
            reference_uri=candidate.get("document_uri"),
            confidence=0.8,
        )
    # Phrase route fired but no clear reference-work source — the
    # verbatim phrase appeared in non-fiction context. Could be a
    # quotation in a history article, a proverbs page, etc.
    return FrameDetection(frame_kind="ambiguous", confidence=0.4)