333 lines
12 KiB
Python
333 lines
12 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Tests for the multi-pass speaker-attribution pipeline.
|
|
|
|
Covers the mechanical narration-evidence builder, speaker normalization for
|
|
parenthetical qualifiers, target collection for the verification pass, and the
|
|
deterministic-fix + LLM verification loop itself (with a stubbed LLM).
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from umbraaudio.analysis.llm_attributor import LLMAttributor
|
|
|
|
|
|
def _attributor_with_speakers(speakers: list[dict]) -> LLMAttributor:
|
|
attr = LLMAttributor.__new__(LLMAttributor)
|
|
attr._load_speakers = lambda: speakers # type: ignore[method-assign]
|
|
return attr
|
|
|
|
|
|
_SPEAKERS = [
|
|
{
|
|
"primary_name": "Kainda Walters Foster Murphy",
|
|
"aliases": ["Princess Kainda", "Princess", "Kainda"],
|
|
},
|
|
{
|
|
"primary_name": "Narrator",
|
|
"aliases": ["narrator", "the narrator"],
|
|
},
|
|
{
|
|
"primary_name": "Chief of the Guard",
|
|
"aliases": ["Lead Guard"],
|
|
},
|
|
{
|
|
"primary_name": "Nanny",
|
|
"aliases": ["Niñera"],
|
|
},
|
|
{
|
|
"primary_name": "Deal Hunter",
|
|
"aliases": ["The AI", "Ship AI"],
|
|
},
|
|
{
|
|
"primary_name": "Xylia",
|
|
"aliases": [],
|
|
},
|
|
]
|
|
|
|
_NARRATION = "narration"
|
|
_DIALOGUE = "dialogue"
|
|
|
|
|
|
def _utt(idx: int, text: str, utt_type: str = _NARRATION, speaker: str | None = None) -> dict:
|
|
u: dict = {
|
|
"utterance_index": idx,
|
|
"text": text,
|
|
"type": utt_type,
|
|
}
|
|
if speaker is not None:
|
|
u["speaker_name"] = speaker
|
|
return u
|
|
|
|
|
|
class TestMechanicalEvidence:
|
|
def test_name_tag_from_preceding_narration(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
scene = [
|
|
_utt(1, "Kainda asked quietly."),
|
|
_utt(2, "Have you seen the vault?", _DIALOGUE),
|
|
]
|
|
evidence = attr._build_mechanical_evidence(scene, [scene[1]], alias_map)
|
|
entry = evidence[2]
|
|
assert entry["tag"] == "Kainda Walters Foster Murphy"
|
|
assert entry["tag_kind"] == "name"
|
|
assert "smiled" in entry["narr_before"] or "asked" in entry["narr_before"]
|
|
|
|
def test_role_tag(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
scene = [
|
|
_utt(1, "The lead guard replied."),
|
|
_utt(2, "Nobody enters tonight.", _DIALOGUE),
|
|
]
|
|
evidence = attr._build_mechanical_evidence(scene, [scene[1]], alias_map)
|
|
entry = evidence[2]
|
|
assert entry["tag_kind"] == "role"
|
|
assert entry["tag"].lower() == "the lead guard"
|
|
|
|
def test_pronoun_tag(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
scene = [
|
|
_utt(1, "She asks, looking nervous."),
|
|
_utt(2, "Where is Father?", _DIALOGUE),
|
|
]
|
|
evidence = attr._build_mechanical_evidence(scene, [scene[1]], alias_map)
|
|
entry = evidence[2]
|
|
assert entry["tag_kind"] == "pronoun"
|
|
assert entry["tag"] == "she"
|
|
|
|
def test_addressee_detected(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
scene = [
|
|
_utt(1, "Nanny stepped closer."),
|
|
_utt(2, "Princess, you must run.", _DIALOGUE),
|
|
]
|
|
evidence = attr._build_mechanical_evidence(scene, [scene[1]], alias_map)
|
|
entry = evidence[2]
|
|
assert entry["addressees"] == ["Kainda Walters Foster Murphy"]
|
|
|
|
def test_no_evidence_when_prev_is_dialogue(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
scene = [
|
|
_utt(1, "First line.", _DIALOGUE, "Kainda Walters Foster Murphy"),
|
|
_utt(2, "Second line.", _DIALOGUE),
|
|
]
|
|
evidence = attr._build_mechanical_evidence(scene, [scene[1]], alias_map)
|
|
entry = evidence[2]
|
|
assert entry["tag"] is None
|
|
assert entry["narr_before"] == ""
|
|
|
|
|
|
class TestNormalizeSpeakerParenthetical:
|
|
def test_strips_parenthetical_qualifier_to_known_name(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
assert (
|
|
attr._normalize_speaker("Xylia (Voice (Dragon))", alias_map)
|
|
== "Xylia"
|
|
)
|
|
|
|
def test_strips_scene_qualifier_to_known_name(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
assert (
|
|
attr._normalize_speaker("Nanny (Scene 14)", alias_map) == "Nanny"
|
|
)
|
|
|
|
def test_keeps_unknown_name_with_parens_unchanged(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
assert attr._normalize_speaker("Somebody (Guard)", alias_map) == "Somebody (Guard)"
|
|
|
|
def test_strips_quotes(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
assert attr._normalize_speaker('"Kainda"', alias_map) == "Kainda Walters Foster Murphy"
|
|
|
|
|
|
class TestCollectVerifyTargets:
|
|
def test_flags_narrator_left_dialogue(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
utterances = [
|
|
_utt(1, "Nanny speaks."),
|
|
_utt(2, "Line one.", _DIALOGUE, "narrator"),
|
|
]
|
|
targets = attr._collect_verify_targets(utterances, 6, alias_map)
|
|
assert 2 in targets
|
|
|
|
def test_flags_tag_mismatch(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
utterances = [
|
|
_utt(1, "Nanny asked"),
|
|
_utt(2, "Line one.", _DIALOGUE, "Kainda Walters Foster Murphy"),
|
|
]
|
|
targets = attr._collect_verify_targets(utterances, 6, alias_map)
|
|
assert 2 in targets
|
|
assert "speech_tag_mismatch" in targets[2]
|
|
|
|
def test_flags_low_confidence(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
utterances = [
|
|
_utt(1, "Some narration."),
|
|
_utt(2, "Confident line.", _DIALOGUE, "Nanny"),
|
|
_utt(3, "Low confidence.", _DIALOGUE, "Nanny"),
|
|
]
|
|
utterances[1]["attribution_confidence"] = 0.75
|
|
utterances[2]["attribution_confidence"] = 0.55
|
|
targets = attr._collect_verify_targets(utterances, 6, alias_map)
|
|
assert 3 in targets # 0.55 < 0.6 → flagged
|
|
assert 2 not in targets # 0.75 ≥ 0.6 → clean
|
|
|
|
def test_flags_uncertain_flag(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
utterances = [
|
|
_utt(1, "Some narration."),
|
|
_utt(2, "Line one.", _DIALOGUE, "Nanny"),
|
|
]
|
|
utterances[1]["attribution_confidence"] = 0.8
|
|
utterances[1]["attribution_uncertain"] = True
|
|
targets = attr._collect_verify_targets(utterances, 6, alias_map)
|
|
assert 2 in targets
|
|
assert "uncertain" in targets[2]
|
|
|
|
|
|
class TestVerifyAttributions:
|
|
def test_applies_llm_correction(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
|
|
def fake_llm(prompt: str) -> list[dict]:
|
|
return [
|
|
{
|
|
"index": 4,
|
|
"speaker": "Nanny",
|
|
"confidence": 0.95,
|
|
"reasoning": "Nanny is the nurse",
|
|
}
|
|
]
|
|
|
|
attr._call_llm = fake_llm # type: ignore[method-assign]
|
|
|
|
utterances = [
|
|
_utt(1, "Nanny entered the room."),
|
|
_utt(2, "Who are you?", _DIALOGUE, "Kainda Walters Foster Murphy"),
|
|
_utt(3, "The nurse."),
|
|
_utt(4, "I am here to help you.", _DIALOGUE, "Kainda Walters Foster Murphy"),
|
|
]
|
|
profiles = "- Kainda Walters Foster Murphy\n- Nanny"
|
|
result = attr._verify_attributions(
|
|
utterances, 6, "Chapter Six", profiles, alias_map, max_rounds=1
|
|
)
|
|
assert result[3]["speaker_name"] == "Nanny"
|
|
|
|
def test_keeps_correct_attribution_unchanged(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
|
|
def fake_llm(prompt: str) -> list[dict]:
|
|
return [{"index": 2, "speaker": "Nanny", "confidence": 0.9, "reasoning": ""}]
|
|
|
|
attr._call_llm = fake_llm # type: ignore[method-assign]
|
|
|
|
utterances = [
|
|
_utt(1, "Nanny spoke."),
|
|
_utt(2, "Stay still.", _DIALOGUE, "Nanny"),
|
|
]
|
|
profiles = "- Nanny"
|
|
result = attr._verify_attributions(
|
|
utterances, 6, "Chapter Six", profiles, alias_map, max_rounds=1
|
|
)
|
|
assert result[1]["speaker_name"] == "Nanny"
|
|
|
|
def test_rejects_narrator_correction(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
|
|
def fake_llm(prompt: str) -> list[dict]:
|
|
return [{"index": 2, "speaker": "Narrator", "confidence": 0.9, "reasoning": ""}]
|
|
|
|
attr._call_llm = fake_llm # type: ignore[method-assign]
|
|
|
|
utterances = [
|
|
_utt(1, "Nanny spoke."),
|
|
_utt(2, "Stay still.", _DIALOGUE, "Nanny"),
|
|
]
|
|
profiles = "- Nanny"
|
|
result = attr._verify_attributions(
|
|
utterances, 6, "Chapter Six", profiles, alias_map, max_rounds=1
|
|
)
|
|
assert result[1]["speaker_name"] == "Nanny"
|
|
|
|
def test_deterministic_fix_applied_before_llm(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
|
|
def fake_llm(prompt: str) -> list[dict]:
|
|
return []
|
|
|
|
attr._call_llm = fake_llm # type: ignore[method-assign]
|
|
|
|
# "Nanny asked" is directly before the line → fix_high_confidence corrects.
|
|
utterances = [
|
|
_utt(1, "Nanny asked"),
|
|
_utt(2, "Who is there?", _DIALOGUE, "Kainda Walters Foster Murphy"),
|
|
]
|
|
profiles = "- Nanny\n- Kainda Walters Foster Murphy"
|
|
result = attr._verify_attributions(
|
|
utterances, 6, "Chapter Six", profiles, alias_map, max_rounds=1
|
|
)
|
|
assert result[1]["speaker_name"] == "Nanny"
|
|
|
|
def test_rejects_invented_name(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
|
|
def fake_llm(prompt: str) -> list[dict]:
|
|
# Digits make it fail both the alias map and the role-name heuristic.
|
|
return [{"index": 2, "speaker": "Guard 5", "confidence": 0.9, "reasoning": ""}]
|
|
|
|
attr._call_llm = fake_llm # type: ignore[method-assign]
|
|
|
|
utterances = [
|
|
_utt(1, "Nanny spoke."),
|
|
_utt(2, "Stay still.", _DIALOGUE, "Nanny"),
|
|
]
|
|
utterances[1]["attribution_confidence"] = 0.5
|
|
profiles = "- Nanny"
|
|
result = attr._verify_attributions(
|
|
utterances, 6, "Chapter Six", profiles, alias_map, max_rounds=1
|
|
)
|
|
assert result[1]["speaker_name"] == "Nanny"
|
|
|
|
def test_verify_targets_batched_at_twenty(self) -> None:
|
|
attr = _attributor_with_speakers(_SPEAKERS)
|
|
alias_map = attr._build_alias_map()
|
|
calls: list[str] = []
|
|
|
|
def fake_llm(prompt: str) -> list[dict]:
|
|
calls.append(prompt)
|
|
return []
|
|
|
|
attr._call_llm = fake_llm # type: ignore[method-assign]
|
|
|
|
# 25 narrator-left dialogue lines → 25 targets → 20 + 5 in two calls.
|
|
utterances = [_utt(1, "Some narration.")]
|
|
for i in range(2, 27):
|
|
utterances.append(_utt(i, f"Line {i}.", _DIALOGUE, "narrator"))
|
|
profiles = "- Nanny"
|
|
attr._verify_attributions(
|
|
utterances, 6, "Chapter Six", profiles, alias_map, max_rounds=1
|
|
)
|
|
assert len(calls) == 2
|
|
assert '"index": 2' in calls[0]
|
|
assert '"index": 22' not in calls[0]
|
|
assert '"index": 22' in calls[1]
|