UmbraAudio/tests/test_attribution_qa.py

from __future__ import annotations

from umbraaudio.analysis.attribution_qa import (
    check_attribution,
    fix_high_confidence,
)

ALIAS_MAP = {
    "hunter": "Deal Hunter",
    "deal hunter": "Deal Hunter",
    "kainda": "Kainda Walters Foster Murphy",
    "kainda walters foster murphy": "Kainda Walters Foster Murphy",
    "princess kainda": "Kainda Walters Foster Murphy",
    "kabir": "Kabir",
    "narrator": "Narrator",
}


def _utt(idx: int, text: str, speaker: str, type_: str = "narration") -> dict:
    return {
        "utterance_index": idx,
        "text": text,
        "type": type_,
        "speaker_name": speaker,
    }


class TestSpeechTagMismatch:
    def test_correct_attribution_no_issue(self) -> None:
        utts = [
            _utt(0, "Hunter says,", "narrator"),
            _utt(1, "We are ready to jump.", "Deal Hunter", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        assert result.has("speech_tag_mismatch") is False

    def test_wrong_attribution_flagged(self) -> None:
        utts = [
            _utt(0, "Hunter says,", "narrator"),
            _utt(1, "We are ready to jump.", "Kabir", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        issues = result.by_type("speech_tag_mismatch")
        assert len(issues) == 1
        assert issues[0].expected == "Deal Hunter"
        assert issues[0].utterance_index == 1
        assert issues[0].confidence >= 0.9

    def test_alias_resolution(self) -> None:
        utts = [
            _utt(0, "Princess Kainda asks,", "narrator"),
            _utt(1, "Where are we?", "Deal Hunter", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        issues = result.by_type("speech_tag_mismatch")
        assert issues[0].expected == "Kainda Walters Foster Murphy"

    def test_tag_not_carried_across_unrelated_narration(self) -> None:
        utts = [
            _utt(0, "Hunter says,", "narrator"),
            _utt(1, "The ship shudders.", "narrator"),
            _utt(2, "We are ready to jump.", "Kabir", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        assert result.has("speech_tag_mismatch") is False

    def test_tag_not_carried_across_scene_break(self) -> None:
        utts = [
            _utt(0, "Hunter says,", "narrator"),
            _utt(1, "<<<>>>", "scene_break"),
            _utt(2, "We are ready to jump.", "Kabir", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        assert result.has("speech_tag_mismatch") is False


class TestNarratorDialogue:
    def test_dialogue_left_as_narrator_flagged(self) -> None:
        utts = [
            _utt(0, "She says,", "narrator"),
            _utt(1, "Hello.", "narrator", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        issues = result.by_type("narrator_dialogue")
        assert len(issues) == 1
        assert issues[0].utterance_index == 1


class TestRoleTag:
    def test_role_tag_assigned_to_named_character_flagged(self) -> None:
        utts = [
            _utt(0, "The lead guard replies,", "narrator"),
            _utt(1, "We have strict orders to keep you away.", "Kabir", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        issues = result.by_type("role_tag_guessed")
        assert len(issues) == 1
        assert issues[0].utterance_index == 1

    def test_role_tag_assigned_to_role_speaker_ok(self) -> None:
        utts = [
            _utt(0, "The lead guard replies,", "narrator"),
            _utt(1, "We have strict orders to keep you away.", "Chief of the Guard", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        assert result.has("role_tag_guessed") is False

    def test_pronoun_tag_skipped(self) -> None:
        utts = [
            _utt(0, "He replies,", "narrator"),
            _utt(1, "We have strict orders to keep you away.", "Chief of the Guard", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        assert not result.issues

    def test_tag_not_applied_to_shared_tag_block(self) -> None:
        utts = [
            _utt(0, "Kainda asks,", "narrator"),
            _utt(
                1, "Where is the ship that brought me?", "Kainda Walters Foster Murphy", "dialogue"
            ),
            _utt(
                2,
                "The ship jumped away as soon as you were on your way.",
                "Willamina Wake",
                "dialogue",
            ),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        assert not result.issues

    def test_voice_tag_assigned_to_ship_flagged(self) -> None:
        utts = [
            _utt(0, "The voice calls again,", "narrator"),
            _utt(1, "Help me, please.", "Deal Hunter", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        issues = result.by_type("role_tag_guessed")
        assert len(issues) == 1

    def test_unresolved_title_word_overlap_ok(self) -> None:
        utts = [
            _utt(0, "Wakam says,", "narrator"),
            _utt(1, "You work a tough deal.", "Trader Wakam", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        assert not result.issues

    def test_unresolved_title_no_overlap_flagged_low_conf(self) -> None:
        utts = [
            _utt(0, "Wakam says,", "narrator"),
            _utt(1, "You work a tough deal.", "Captain Weston", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        issues = result.by_type("speech_tag_mismatch")
        assert len(issues) == 1
        assert issues[0].confidence < 0.9

    def test_conjunction_not_absorbed_into_role(self) -> None:
        utts = [
            _utt(0, "A supervisor sees the robot and tells it to stop.", "narrator"),
            _utt(1, "Why are you here?", "Human Staff Person/Messenger", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        assert not result.issues

    def test_plain_narration_not_flagged(self) -> None:
        utts = [
            _utt(0, "The ship powers down its engines.", "narrator"),
            _utt(1, "Help me, please.", "Deal Hunter", "dialogue"),
        ]
        result = check_attribution(utts, 6, ALIAS_MAP)
        assert not result.issues


class TestFixHighConfidence:
    def test_fixes_exact_name_mismatch_only(self) -> None:
        utts = [
            _utt(0, "Hunter says,", "narrator"),
            _utt(1, "We are ready to jump.", "Kabir", "dialogue"),
            _utt(2, "The ship shudders.", "narrator"),
            _utt(3, "Can you take me on a tour?", "Kabir", "dialogue"),
        ]
        fixed = fix_high_confidence(utts, ALIAS_MAP)
        assert fixed[1]["speaker_name"] == "Deal Hunter"
        assert fixed[3]["speaker_name"] == "Kabir"

    def test_already_correct_unchanged(self) -> None:
        utts = [
            _utt(0, "Hunter says,", "narrator"),
            _utt(1, "We are ready to jump.", "Deal Hunter", "dialogue"),
        ]
        fixed = fix_high_confidence(utts, ALIAS_MAP)
        assert fixed[1]["speaker_name"] == "Deal Hunter"