tests/test_attention_compiler.py

import math

import pytest

from app.schemas.visual_lesson import AttentionLessonDraft, EvidenceClaim, EvidenceSource
from app.services.attention_compiler import AttentionCompiler, AttentionCompilerError, _matmul, _softmax_rows


def _draft(tokens=None):
    return AttentionLessonDraft(
        title="Inside scaled dot-product attention",
        interpretation="A small deterministic walkthrough of self-attention.",
        requested_variant="standard_self_attention",
        lesson_variant="standard_self_attention",
        variant_notice="",
        tokens=tokens or ["Research", "Mate", "reads", "papers"],
        assumptions=["The displayed tensors are illustrative teaching values."],
        teaching_steps=[
            "Start with tokens and embeddings.",
            "Project embeddings into queries, keys, and values.",
            "Compare queries with keys.",
            "Scale the logits.",
            "Optionally apply a causal mask.",
            "Normalize each row with softmax.",
            "Aggregate values into the output.",
        ],
        claim_ids=["canonical-attention", "illustrative-values"],
    )


def _evidence():
    sources = [
        EvidenceSource(
            source_id="builtin:vaswani-2017",
            origin="builtin",
            title="Attention Is All You Need",
            url="https://arxiv.org/abs/1706.03762",
            authority="canonical",
        )
    ]
    claims = [
        EvidenceClaim(
            claim_id="canonical-attention",
            text="Scaled dot-product attention applies softmax to scaled query-key products before aggregating values.",
            claim_type="standard_definition",
            support_level="canonical",
            source_ids=["builtin:vaswani-2017"],
        ),
        EvidenceClaim(
            claim_id="illustrative-values",
            text="The small matrices in this lesson are illustrative rather than trained model weights.",
            claim_type="illustrative_choice",
            support_level="illustrative",
            source_ids=[],
        ),
    ]
    return sources, claims


def test_compiler_produces_verified_unmasked_and_causal_branches():
    sources, claims = _evidence()
    payload = AttentionCompiler().build_payload(
        project_id="project-1",
        prompt="visualize attention",
        draft=_draft(),
        evidence_sources=sources,
        evidence_claims=claims,
        warnings=[],
    )

    assert [branch.branch_id for branch in payload.compiled.branches] == ["unmasked", "causal"]
    assert len(payload.spec.timeline) == 7
    assert payload.compiled.assertions_passed is True

    for branch in payload.compiled.branches:
        for row in branch.tensors["attentionWeights"]:
            assert math.isclose(sum(row), 1.0, abs_tol=1e-6)

    causal = payload.compiled.branches[1].tensors["attentionWeights"]
    for row_index, row in enumerate(causal):
        assert all(abs(value) <= 1e-6 for value in row[row_index + 1 :])


def test_compiler_replaces_unsafe_or_wrong_sized_token_labels():
    sources, claims = _evidence()
    payload = AttentionCompiler().build_payload(
        project_id="project-1",
        prompt="visualize attention",
        draft=_draft(tokens=["only", "three", "tokens"]),
        evidence_sources=sources,
        evidence_claims=claims,
        warnings=[],
    )

    assert payload.spec.tokens == ["Research", "Mate", "reads", "papers"]


def test_compiler_rejects_an_unapproved_operation_graph():
    compiler = AttentionCompiler()
    operations = compiler.operation_graph()
    operations[0] = operations[0].model_copy(update={"op": "identity"})

    with pytest.raises(AttentionCompilerError, match="approved attention operation graph"):
        compiler.compile_operations(operations, causal=False)


def test_matrix_multiplication_and_stable_softmax_are_deterministic():
    assert _matmul([[1.0, 2.0]], [[3.0], [4.0]]) == [[11.0]]
    probabilities = _softmax_rows([[1000.0, 1001.0, 999.0]])[0]
    assert all(math.isfinite(value) for value in probabilities)
    assert math.isclose(sum(probabilities), 1.0, abs_tol=1e-12)
    assert probabilities[1] > probabilities[0] > probabilities[2]


def test_fixed_fixture_recompiles_to_identical_tensors():
    sources, claims = _evidence()
    compiler = AttentionCompiler()
    payload = compiler.build_payload(
        project_id="project-1", prompt="attention", draft=_draft(),
        evidence_sources=sources, evidence_claims=claims, warnings=[],
    )

    recompiled = compiler.compile_spec(payload.spec)
    assert recompiled.model_dump() == payload.compiled.model_dump()


def test_compiler_rejects_nonfinite_or_wrong_shaped_persisted_inputs():
    compiler = AttentionCompiler()
    with pytest.raises(AttentionCompilerError, match="embeddings must have shape"):
        compiler.compile_operations(
            compiler.operation_graph(), causal=False,
            input_tensors={"embeddings": [[1.0]], "Wq": [[1.0, 0.0]] * 3, "Wk": [[1.0, 0.0]] * 3, "Wv": [[1.0, 0.0]] * 3},
        )