tests/test_modality_router.py

"""Tests for VisualModalityRouter -> verifies routing rules and prompt grounding."""
from unittest.mock import MagicMock

from app.agents.modality_router import VisualModalityRouter


def _make_router_with_response(modality: str, decline_reason: str = "") -> tuple:
    """Return (router, captured_messages) where the client returns a canned decision."""
    router = VisualModalityRouter.__new__(VisualModalityRouter)
    mock_client = MagicMock()
    router._client = mock_client

    from app.agents.modality_router import ModalityDecision
    decision = ModalityDecision(
        modality=modality,
        reasoning="test reasoning",
        decline_reason=decline_reason,
    )
    captured = {}

    def fake_structured_complete(messages, model_cls):
        captured["messages"] = messages
        return decision

    mock_client.structured_complete.side_effect = fake_structured_complete
    return router, captured


CHUNKS_WITH_DATA = [
    {"source": "notes.pdf", "text": "Revenue in 2020: $120M, 2021: $145M, 2022: $178M"},
    {"source": "notes.pdf", "text": "Growth rate was 20.8% year over year."},
]

CHUNKS_WITHOUT_DATA = [
    {"source": "notes.pdf", "text": "Entropy is a measure of disorder in a thermodynamic system."},
    {"source": "notes.pdf", "text": "The second law states entropy increases in isolated systems."},
]


def test_classify_returns_modality_decision():
    router, _ = _make_router_with_response("graph")
    result = router.classify("revenue trends", "## Revenue\n...", CHUNKS_WITH_DATA, "high_school")
    assert result.modality == "graph"


def test_prompt_contains_chunk_text():
    """The router must embed source material in the prompt so decisions are grounded."""
    router, captured = _make_router_with_response("graph")
    router.classify("revenue", "card text", CHUNKS_WITH_DATA, "graduate")
    user_message = captured["messages"][1]["content"]
    assert "Revenue in 2020" in user_message, "chunk text must appear in the user message"
    assert "notes.pdf" not in user_message  # chunk source key is not embedded, just the text


def test_prompt_contains_selection_and_level():
    router, captured = _make_router_with_response("decline", decline_reason="purely definitional")
    router.classify("entropy", "card", CHUNKS_WITHOUT_DATA, "eli5")
    user_message = captured["messages"][1]["content"]
    assert "entropy" in user_message
    assert "eli5" in user_message


def test_prompt_caps_chunk_text():
    """Chunks larger than 3000 chars must be truncated."""
    big_chunk = [{"source": "big.pdf", "text": "x" * 5000}]
    router, captured = _make_router_with_response("decline", decline_reason="no data")
    router.classify("concept", "card", big_chunk, "high_school")
    user_message = captured["messages"][1]["content"]
    # The chunk text is capped at 3000 chars inside classify()
    assert len(user_message) < 5000 + 500  # well under 5000 raw chunk length


def test_decline_modality_when_no_data():
    router, _ = _make_router_with_response("decline", decline_reason="purely conceptual, no data or dynamic process")
    result = router.classify("entropy definition", "card", CHUNKS_WITHOUT_DATA, "high_school")
    assert result.modality == "decline"
    assert result.decline_reason


def test_2d_anim_for_dynamic_concept():
    router, _ = _make_router_with_response("2d_anim")
    result = router.classify("pendulum motion", "card", CHUNKS_WITHOUT_DATA, "high_school")
    assert result.modality == "2d_anim"


def test_all_six_modalities_are_valid_literal_values():
    """Every value in the new 6-way taxonomy must be constructible on ModalityDecision."""
    from app.agents.modality_router import ModalityDecision
    for modality in ["formula", "graph", "2d_text", "3d", "2d_anim", "decline"]:
        decision = ModalityDecision(modality=modality, reasoning="r")
        assert decision.modality == modality


def test_decline_is_standalone_outcome_not_and_gated():
    """Regression test for the confirmed router bug: the old prompt only reached NONE/decline
    when BOTH 'no data' AND 'no dynamic process' held simultaneously (an AND-gate so narrow
    that almost nothing reached it, defaulting to offering a visual almost every time).

    The rewritten prompt must present `decline` as a directly reachable, standalone
    classification outcome alongside the other 5 modalities -- not a fallback gated behind
    a conjunction of two other conditions.
    """
    from app.agents.modality_router import _SYSTEM_PROMPT

    prompt_lower = _SYSTEM_PROMPT.lower()

    # The old AND-gated phrasing must not reappear in any form.
    banned_phrases = [
        "no data and no dynamic",
        "no data or dynamic system",  # old NONE rule wording (still a conjunction of two conditions)
        "must not choose",
    ]
    for phrase in banned_phrases:
        assert phrase not in prompt_lower, f"AND-gated decline phrasing reintroduced: {phrase!r}"

    # `decline` must be listed as one of the parallel, independently-selectable modalities,
    # not merely mentioned in passing.
    assert "decline" in prompt_lower
    # All 6 outcomes should appear as options the model picks ONE of.
    for modality in ["formula", "graph", "2d_text", "3d", "2d_anim", "decline"]:
        assert modality in prompt_lower, f"modality {modality!r} missing from system prompt"

    # Positive check: banning old phrasing isn't enough on its own -- a future edit could
    # reintroduce the same AND-gate bug in different words while still avoiding the banned
    # substrings above and mentioning all 6 modality names. So also require the actual
    # parallel/standalone/independent framing language to be present verbatim.
    standalone_framing_phrases = [
        "independent, parallel options",  # states all 6 outcomes are independent, not chained
        "equally reachable as a direct, standalone choice",  # decline is not a rare fallback
        "just as valid an answer as any other modality",  # decline needs no other options ruled out first
    ]
    for phrase in standalone_framing_phrases:
        assert phrase in prompt_lower, f"expected standalone-framing language missing: {phrase!r}"