arborist/tests/test_frame.py

"""Frame-detection tests (#000002 / Module L scope).

Pins the heuristic detector that classifies a claim-lattice query as
``literal`` / ``reference`` / ``ambiguous`` / ``no_phrase_route`` based
on the phrase-pattern retrieval signal + per-source reference-work
indicators.
"""
from __future__ import annotations

from arborist.qa.frame import (
    FrameDetection,
    _body_indicates_reference_work,
    _source_is_reference_work,
    _title_indicates_reference_work,
    detect_frame,
)


# ---------------------------------------------------------------------------
# Title-suffix detection
# ---------------------------------------------------------------------------


def test_title_suffix_recognises_novel_film_play():
    assert _title_indicates_reference_work("Jurassic Park (film)")
    assert _title_indicates_reference_work("Jurassic Park (novel)")
    assert _title_indicates_reference_work("Hamlet (play)")
    assert _title_indicates_reference_work("The Lord of the Rings (franchise)")
    assert _title_indicates_reference_work("Star Wars (video game)")


def test_title_suffix_handles_multiword_disambig():
    assert _title_indicates_reference_work("Some Comic (graphic novel)")
    assert _title_indicates_reference_work("Some Show (TV series)")
    assert _title_indicates_reference_work("Some Plot (short story)")


def test_title_suffix_doesnt_fire_on_literal_geography():
    assert not _title_indicates_reference_work("Oceania")
    assert not _title_indicates_reference_work("Asia")
    assert not _title_indicates_reference_work("New York City")
    assert not _title_indicates_reference_work("Springfield, Missouri")


def test_title_suffix_doesnt_fire_on_disambig_for_real_things():
    """Geographic / biographical disambig suffixes should NOT trip
    the reference-work detector."""
    assert not _title_indicates_reference_work("Cleveland (Ohio)")
    assert not _title_indicates_reference_work("Mercury (planet)")
    assert not _title_indicates_reference_work("Mercury (element)")
    assert not _title_indicates_reference_work(
        "John Smith (politician)"
    )


def test_title_suffix_handles_none_or_empty():
    assert not _title_indicates_reference_work(None)
    assert not _title_indicates_reference_work("")


# ---------------------------------------------------------------------------
# Body fiction-marker density
# ---------------------------------------------------------------------------


def test_body_density_recognises_orwell_lead():
    """The Nineteen_Eighty-Four article opens with high fiction-
    marker density: 'novel by George Orwell' + 'fiction' + 'plot' /
    'characters' / 'published' typically all in the first KB."""
    sample = (
        "Nineteen Eighty-Four is a dystopian social science fiction "
        "novel by English novelist George Orwell. The narrative "
        "follows the protagonist Winston Smith, a low-ranking "
        "member of the ruling Party. The novel was published in "
        "1949. The plot explores themes of totalitarianism. "
        "Characters include Big Brother and Julia."
    )
    assert _body_indicates_reference_work(sample)


def test_body_density_doesnt_fire_on_geography():
    sample = (
        "Oceania is a continent located in the Pacific Ocean. It "
        "comprises Australia, New Zealand, and various smaller "
        "Pacific island nations. The continent has a diverse "
        "geography with mountains, deserts, and coral reefs. The "
        "population is concentrated along the coastlines."
    )
    assert not _body_indicates_reference_work(sample)


def test_body_density_doesnt_fire_on_one_incidental_novel_use():
    """A history article saying 'this was a novel approach' shouldn't
    trip the detector — single use vs marker density."""
    sample = (
        "The reform of 1850 introduced a novel approach to "
        "regional administration. Local governors retained "
        "significant autonomy under the new framework. The "
        "policy persisted for several decades before being "
        "replaced. " * 2
    )
    assert not _body_indicates_reference_work(sample)


def test_body_density_handles_none_or_empty():
    assert not _body_indicates_reference_work(None)
    assert not _body_indicates_reference_work("")


# ---------------------------------------------------------------------------
# _source_is_reference_work — composition
# ---------------------------------------------------------------------------


def test_source_is_reference_work_via_title():
    src = {
        "title": "Hamlet (play)",
        "body_sample": "Hamlet is a tragedy.",
    }
    assert _source_is_reference_work(src)


def test_source_is_reference_work_via_body_only():
    src = {
        "title": "Nineteen Eighty-Four",
        "body_sample": (
            "Nineteen Eighty-Four is a dystopian science fiction "
            "novel by George Orwell, published in 1949. The plot "
            "follows Winston Smith. The novel explores themes of "
            "totalitarianism. Characters include Big Brother."
        ),
    }
    assert _source_is_reference_work(src)


def test_source_is_not_reference_work_for_geography():
    src = {
        "title": "Oceania",
        "body_sample": (
            "Oceania is a geographic region comprising Australia, "
            "New Zealand, and Pacific Islands."
        ),
    }
    assert not _source_is_reference_work(src)


# ---------------------------------------------------------------------------
# detect_frame — end-to-end classification
# ---------------------------------------------------------------------------


def test_detect_frame_no_phrase_route():
    """No phrase route fired → no_phrase_route classification."""
    sources = [{"document_root": "abc", "title": "Oceania"}]
    fd = detect_frame("oceania population", sources, phrase_match_roots=set())
    assert fd.frame_kind == "no_phrase_route"
    assert fd.reference_title is None


def test_detect_frame_no_phrase_route_when_arg_is_none():
    sources = [{"document_root": "abc", "title": "Oceania"}]
    fd = detect_frame("oceania population", sources, phrase_match_roots=None)
    assert fd.frame_kind == "no_phrase_route"


def test_detect_frame_reference_via_title_suffix():
    """Phrase route fired AND matched source has a (film) suffix →
    reference-frame classification."""
    sources = [
        {
            "document_root": "aaa",
            "title": "Jurassic Park (film)",
            "document_uri": "https://en.wikipedia.org/wiki/Jurassic_Park_(film)",
            "body_sample": "Jurassic Park is a 1993 American science fiction action film.",
        },
    ]
    fd = detect_frame(
        "what dinosaurs were in jurassic park?",
        sources,
        phrase_match_roots={"aaa"},
    )
    assert fd.frame_kind == "reference"
    assert fd.reference_title == "Jurassic Park (film)"
    assert fd.confidence > 0.5


def test_detect_frame_reference_via_body_density_orwell():
    """Phrase route fired AND matched source has fiction body markers
    → reference-frame even without a title suffix (Nineteen
    Eighty-Four canonical case)."""
    sources = [
        {
            "document_root": "bbb",
            "title": "Nineteen Eighty-Four",
            "document_uri": "https://en.wikipedia.org/wiki/Nineteen_Eighty-Four",
            "body_sample": (
                "Nineteen Eighty-Four is a dystopian science fiction "
                "novel by George Orwell, published in 1949. The plot "
                "follows Winston Smith, a Party member. Characters "
                "include Big Brother and Julia. The novel explores "
                "totalitarianism."
            ),
        },
    ]
    fd = detect_frame(
        "has oceania always been at war with east asia",
        sources,
        phrase_match_roots={"bbb"},
    )
    assert fd.frame_kind == "reference"
    assert fd.reference_title == "Nineteen Eighty-Four"


def test_detect_frame_ambiguous_when_phrase_matched_but_not_reference():
    """Phrase route fired but the matched source is a non-fiction
    article (e.g. a verbatim quotation that happens to appear in a
    history page). Classification: ambiguous."""
    sources = [
        {
            "document_root": "ccc",
            "title": "List of common phrases",
            "document_uri": "https://example.org/list",
            "body_sample": (
                "This is a list page. It contains various phrases "
                "and their meanings. Geographic regions include "
                "Asia, Africa, and Europe."
            ),
        },
    ]
    fd = detect_frame(
        "always been at war with east asia",
        sources,
        phrase_match_roots={"ccc"},
    )
    assert fd.frame_kind == "ambiguous"
    assert fd.reference_title is None


def test_detect_frame_picks_first_reference_source_when_multiple():
    """When phrase route surfaced multiple sources and ≥2 are
    reference works, pick the first by retrieval order."""
    sources = [
        {
            "document_root": "aaa",
            "title": "Jurassic Park (film)",
            "body_sample": "An action film.",
        },
        {
            "document_root": "bbb",
            "title": "Jurassic Park (novel)",
            "body_sample": "A novel by Michael Crichton.",
        },
    ]
    fd = detect_frame(
        "dinosaurs in jurassic park",
        sources,
        phrase_match_roots={"aaa", "bbb"},
    )
    assert fd.frame_kind == "reference"
    # First-by-order is the film.
    assert fd.reference_title == "Jurassic Park (film)"


def test_detect_frame_skips_non_phrase_matched_sources():
    """A reference-work source that wasn't phrase-matched doesn't
    trigger the reference frame (otherwise any retrieval that
    happens to include a (film)-suffixed source would over-fire)."""
    sources = [
        {
            "document_root": "aaa",
            "title": "Hamlet (play)",  # reference work but not phrase-matched
            "body_sample": "A tragedy.",
        },
    ]
    fd = detect_frame(
        "hamlet plot summary",
        sources,
        phrase_match_roots=set(),  # no phrase route
    )
    assert fd.frame_kind == "no_phrase_route"