Files
audita/tests/test_pipeline.py

668 lines
20 KiB
Python

import json
from audita.config import AuditaConfig
from audita.pipeline import process_transcript
from audita.schemas import CorrectionCandidate, CorrectionSet, parse_glossary_yaml, parse_source_transcript_json
class FakeLLMClient:
def __init__(self, responses):
self.responses = list(responses)
self.calls = 0
self.messages = []
def create_corrections(self, messages, config):
self.calls += 1
self.messages.append(messages)
return self.responses.pop(0)
def _config(tmp_path, glossary_max_llm_passes=3, grammar_max_llm_passes=3):
return AuditaConfig(
api_key="key",
max_section_tokens=16000,
glossary_confidence_threshold=0.8,
grammar_confidence_threshold=0.8,
max_retries=3,
glossary_max_llm_passes=glossary_max_llm_passes,
grammar_max_llm_passes=grammar_max_llm_passes,
work_dir=tmp_path / "work",
)
def _glossary():
return parse_glossary_yaml(
"""
glossary:
- name: "Chauntea"
category: deity
summary: "Chauntea is a deity."
"""
)
def _transcript():
return parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "I ask Chontia."},
{"speaker": "Mike", "start": 10.0, "end": 11.0, "text": "Then Lyra."}
]
"""
)
def test_pipeline_processes_with_fake_llm_and_cleans_work_dir(tmp_path):
correction = CorrectionCandidate(
id=1,
original_text="Chontia",
corrected_text="Chauntea",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[correction]),
CorrectionSet(corrections=[]),
]
)
revised = process_transcript(
_transcript(),
_glossary(),
_config(tmp_path),
llm_client=fake_client,
)
assert [segment.speaker for segment in revised] == ["Eric", "Mike"]
assert revised[0].text == "I ask Chauntea."
assert fake_client.calls == 2
assert list((tmp_path / "work").iterdir()) == []
def test_glossary_stage_can_correct_toward_protected_term(tmp_path):
transcript = parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "Frank moves."}
]
"""
)
glossary = parse_glossary_yaml(
"""
glossary:
- name: "Hrank"
category: pc
summary: "Hrank is a player character."
"""
)
glossary_correction = CorrectionCandidate(
id=1,
original_text="Frank",
corrected_text="Hrank",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[glossary_correction]),
CorrectionSet(corrections=[]),
]
)
revised = process_transcript(
transcript,
glossary,
_config(tmp_path),
llm_client=fake_client,
)
assert revised[0].text == "Hrank moves."
assert list((tmp_path / "work").iterdir()) == []
def test_glossary_stage_cannot_change_away_from_protected_term(tmp_path):
transcript = parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "Hrank moves."}
]
"""
)
glossary = parse_glossary_yaml(
"""
glossary:
- name: "Hrank"
category: pc
summary: "Hrank is a player character."
"""
)
glossary_reversal = CorrectionCandidate(
id=1,
original_text="Hrank",
corrected_text="Frank",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[glossary_reversal]),
CorrectionSet(corrections=[]),
]
)
revised = process_transcript(
transcript,
glossary,
_config(tmp_path, glossary_max_llm_passes=1),
llm_client=fake_client,
)
assert revised[0].text == "Hrank moves."
run_dirs = list((tmp_path / "work").iterdir())
assert len(run_dirs) == 1
diagnostics = json.loads((run_dirs[0] / "skipped-corrections.json").read_text(encoding="utf-8"))
skipped = diagnostics["skipped_corrections"][0]
assert skipped["stage"] == "glossary"
assert skipped["reason"] == "correction changes protected glossary term usage"
def test_glossary_stage_cannot_decanonicalize_protected_term(tmp_path):
transcript = parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "Hrank moves."}
]
"""
)
glossary = parse_glossary_yaml(
"""
glossary:
- name: "Hrank"
category: pc
summary: "Hrank is a player character."
"""
)
decapitalization = CorrectionCandidate(
id=1,
original_text="Hrank",
corrected_text="hrank",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[decapitalization]),
CorrectionSet(corrections=[]),
]
)
revised = process_transcript(
transcript,
glossary,
_config(tmp_path, glossary_max_llm_passes=1),
llm_client=fake_client,
)
assert revised[0].text == "Hrank moves."
run_dirs = list((tmp_path / "work").iterdir())
assert len(run_dirs) == 1
diagnostics = json.loads((run_dirs[0] / "skipped-corrections.json").read_text(encoding="utf-8"))
skipped = diagnostics["skipped_corrections"][0]
assert skipped["stage"] == "glossary"
assert skipped["reason"] == "correction changes protected glossary term capitalization"
def test_pipeline_normalizes_before_llm_prompts(tmp_path):
transcript = parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "I ask"},
{"speaker": "Eric", "start": 2.0, "end": 3.0, "text": "Chontia."},
{"speaker": "Mike", "start": 10.0, "end": 11.0, "text": "Then Lyra."}
]
"""
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[]),
CorrectionSet(corrections=[]),
]
)
progress = []
process_transcript(
transcript,
_glossary(),
_config(tmp_path),
llm_client=fake_client,
progress=progress.append,
)
glossary_prompt = fake_client.messages[0][1]["content"]
glossary_payload = json.loads(glossary_prompt.split("Transcript section:\n", maxsplit=1)[1])
assert glossary_payload[0] == {"id": 1, "original_text": "I ask Chontia."}
assert any("Normalized transcript from 3 to 2 segments" in message for message in progress)
def test_pipeline_skips_bad_glossary_correction_and_preserves_diagnostics(tmp_path):
correction = CorrectionCandidate(
id=1,
original_text="Different text.",
corrected_text="Chauntea",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[correction]),
CorrectionSet(corrections=[]),
]
)
progress = []
revised = process_transcript(
_transcript(),
_glossary(),
_config(tmp_path, glossary_max_llm_passes=1),
llm_client=fake_client,
progress=progress.append,
)
assert revised[0].text == "I ask Chontia."
assert any("Skipping glossary correction for id 1" in message for message in progress)
preserved = list((tmp_path / "work").iterdir())
assert len(preserved) == 1
skipped_path = preserved[0] / "skipped-corrections.json"
assert skipped_path.exists()
diagnostics = json.loads(skipped_path.read_text(encoding="utf-8"))
assert diagnostics["skipped_corrections"][0]["stage"] == "glossary"
assert diagnostics["skipped_corrections"][0]["id"] == 1
assert "does not match any substring" in diagnostics["skipped_corrections"][0]["reason"]
def test_pipeline_retries_skipped_segment_and_cleans_work_dir_when_fixed(tmp_path):
first_pass = CorrectionCandidate(
id=1,
original_text="Contia",
corrected_text="Chauntea",
confidence=0.95,
)
second_pass = CorrectionCandidate(
id=1,
original_text="Chontia",
corrected_text="Chauntea",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[first_pass]),
CorrectionSet(corrections=[second_pass]),
CorrectionSet(corrections=[]),
]
)
revised = process_transcript(
_transcript(),
_glossary(),
_config(tmp_path, glossary_max_llm_passes=3),
llm_client=fake_client,
)
assert fake_client.calls == 3
assert [segment.speaker for segment in revised] == ["Eric", "Mike"]
assert revised[0].text == "I ask Chauntea."
assert list((tmp_path / "work").iterdir()) == []
def test_pipeline_retry_prompt_contains_only_valid_deduped_ids(tmp_path):
first_bad = CorrectionCandidate(
id=1,
original_text="Contia",
corrected_text="Chauntea",
confidence=0.95,
)
second_bad_same_segment = CorrectionCandidate(
id=1,
original_text="Still wrong",
corrected_text="Chauntea",
confidence=0.95,
)
invalid_segment = CorrectionCandidate(
id=99,
original_text="Missing",
corrected_text="Chauntea",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[first_bad, second_bad_same_segment, invalid_segment]),
CorrectionSet(corrections=[]),
CorrectionSet(corrections=[]),
]
)
process_transcript(
_transcript(),
_glossary(),
_config(tmp_path, glossary_max_llm_passes=2),
llm_client=fake_client,
)
assert fake_client.calls == 3
retry_prompt = fake_client.messages[1][1]["content"]
retry_payload = json.loads(retry_prompt.split("Transcript section:\n", maxsplit=1)[1])
assert retry_payload == [{"id": 1, "original_text": "I ask Chontia."}]
assert "Retry guidance" in retry_prompt
def test_pipeline_writes_stage_metadata_for_unresolved_retries(tmp_path):
first_pass = CorrectionCandidate(
id=1,
original_text="Contia",
corrected_text="Chauntea",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[first_pass]),
CorrectionSet(corrections=[]),
CorrectionSet(corrections=[]),
]
)
process_transcript(
_transcript(),
_glossary(),
_config(tmp_path, glossary_max_llm_passes=2),
llm_client=fake_client,
)
run_dirs = list((tmp_path / "work").iterdir())
assert len(run_dirs) == 1
assert (run_dirs[0] / "glossary" / "pass-0001").exists()
assert (run_dirs[0] / "grammar" / "pass-0001").exists()
assert (run_dirs[0] / "normalization" / "source-transcript.json").exists()
assert (run_dirs[0] / "normalization" / "normalized-transcript.json").exists()
assert (run_dirs[0] / "normalization" / "summary.json").exists()
metadata = json.loads((run_dirs[0] / "metadata.json").read_text(encoding="utf-8"))
assert metadata["normalization"]["source_segment_count"] == 2
assert metadata["normalization"]["normalized_segment_count"] == 2
assert metadata["normalization"]["merge_count"] == 0
assert metadata["glossary_max_llm_passes"] == 2
assert metadata["grammar_max_llm_passes"] == 3
assert metadata["glossary_confidence_threshold"] == 0.8
assert metadata["grammar_confidence_threshold"] == 0.8
assert [item["stage"] for item in metadata["stages"]] == ["glossary", "grammar"]
assert [item["pass_number"] for item in metadata["stages"][0]["passes"]] == [1, 2]
assert metadata["stages"][0]["passes"][0]["retry_segment_count"] == 1
assert metadata["stages"][0]["passes"][1]["retry_pass"] is True
def test_grammar_stage_runs_after_glossary_and_sees_corrected_text(tmp_path):
transcript = parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "i ask Chontia."},
{"speaker": "Mike", "start": 10.0, "end": 11.0, "text": "Then Lyra."}
]
"""
)
glossary_correction = CorrectionCandidate(
id=1,
original_text="Chontia",
corrected_text="Chauntea",
confidence=0.95,
)
grammar_correction = CorrectionCandidate(
id=1,
original_text="i",
corrected_text="I",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[glossary_correction]),
CorrectionSet(corrections=[grammar_correction]),
]
)
revised = process_transcript(
transcript,
_glossary(),
_config(tmp_path),
llm_client=fake_client,
)
grammar_prompt = fake_client.messages[1][1]["content"]
grammar_payload = json.loads(grammar_prompt.split("Transcript section:\n", maxsplit=1)[1])
assert grammar_payload[0]["original_text"] == "i ask Chauntea."
assert revised[0].text == "I ask Chauntea."
def test_grammar_stage_cannot_reverse_glossary_protected_term(tmp_path):
transcript = parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "Frank moves."}
]
"""
)
glossary = parse_glossary_yaml(
"""
glossary:
- name: "Hrank"
category: pc
summary: "Hrank is a player character."
"""
)
glossary_correction = CorrectionCandidate(
id=1,
original_text="Frank",
corrected_text="Hrank",
confidence=0.95,
)
grammar_reversal = CorrectionCandidate(
id=1,
original_text="Hrank",
corrected_text="Frank",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[glossary_correction]),
CorrectionSet(corrections=[grammar_reversal]),
]
)
progress = []
revised = process_transcript(
transcript,
glossary,
_config(tmp_path, grammar_max_llm_passes=1),
llm_client=fake_client,
progress=progress.append,
)
assert revised[0].text == "Hrank moves."
assert any("Skipping grammar correction for id 1" in message for message in progress)
run_dirs = list((tmp_path / "work").iterdir())
assert len(run_dirs) == 1
diagnostics = json.loads((run_dirs[0] / "skipped-corrections.json").read_text(encoding="utf-8"))
assert diagnostics["skipped_corrections"][0]["stage"] == "grammar"
assert diagnostics["skipped_corrections"][0]["reason"] == "correction changes protected glossary term usage"
def test_grammar_stage_can_correct_toward_protected_term(tmp_path):
transcript = parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "Pawpaw's just worn out."}
]
"""
)
glossary = parse_glossary_yaml(
"""
glossary:
- name: "Popov"
category: npc
summary: "Popov is an allied NPC."
"""
)
grammar_correction = CorrectionCandidate(
id=1,
original_text="Pawpaw's",
corrected_text="Popov's",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[]),
CorrectionSet(corrections=[grammar_correction]),
]
)
revised = process_transcript(
transcript,
glossary,
_config(tmp_path),
llm_client=fake_client,
)
assert revised[0].text == "Popov's just worn out."
assert list((tmp_path / "work").iterdir()) == []
def test_grammar_stage_cannot_change_away_from_protected_term(tmp_path):
transcript = parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "Popov's just worn out."}
]
"""
)
glossary = parse_glossary_yaml(
"""
glossary:
- name: "Popov"
category: npc
summary: "Popov is an allied NPC."
"""
)
grammar_correction = CorrectionCandidate(
id=1,
original_text="Popov's",
corrected_text="Pawpaw's",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[]),
CorrectionSet(corrections=[grammar_correction]),
]
)
revised = process_transcript(
transcript,
glossary,
_config(tmp_path, grammar_max_llm_passes=1),
llm_client=fake_client,
)
assert revised[0].text == "Popov's just worn out."
run_dirs = list((tmp_path / "work").iterdir())
assert len(run_dirs) == 1
diagnostics = json.loads((run_dirs[0] / "skipped-corrections.json").read_text(encoding="utf-8"))
assert diagnostics["skipped_corrections"][0]["stage"] == "grammar"
assert diagnostics["skipped_corrections"][0]["reason"] == "correction changes protected glossary term usage"
def test_grammar_stage_retries_repeated_span_and_applies_unique_retry(tmp_path):
transcript = parse_source_transcript_json(
"""
[
{"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "there and there."}
]
"""
)
repeated_span = CorrectionCandidate(
id=1,
original_text="there",
corrected_text="their",
confidence=0.95,
)
unique_retry = CorrectionCandidate(
id=1,
original_text="there and there",
corrected_text="their and there",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[]),
CorrectionSet(corrections=[repeated_span]),
CorrectionSet(corrections=[unique_retry]),
]
)
revised = process_transcript(
transcript,
_glossary(),
_config(tmp_path, grammar_max_llm_passes=2),
llm_client=fake_client,
)
assert fake_client.calls == 3
assert revised[0].text == "their and there."
retry_prompt = fake_client.messages[2][1]["content"]
retry_payload = json.loads(retry_prompt.split("Transcript section:\n", maxsplit=1)[1])
assert retry_payload == [{"id": 1, "original_text": "there and there."}]
def test_below_threshold_grammar_corrections_are_not_retried(tmp_path):
correction = CorrectionCandidate(
id=1,
original_text="I",
corrected_text="i",
confidence=0.7,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[]),
CorrectionSet(corrections=[correction]),
]
)
revised = process_transcript(
_transcript(),
_glossary(),
_config(tmp_path, grammar_max_llm_passes=3),
llm_client=fake_client,
)
assert fake_client.calls == 2
assert revised[0].text == "I ask Chontia."
def test_unresolved_grammar_skip_preserves_diagnostics(tmp_path):
correction = CorrectionCandidate(
id=1,
original_text="a",
corrected_text="A",
confidence=0.95,
)
fake_client = FakeLLMClient(
[
CorrectionSet(corrections=[]),
CorrectionSet(corrections=[correction]),
]
)
process_transcript(
_transcript(),
_glossary(),
_config(tmp_path, grammar_max_llm_passes=1),
llm_client=fake_client,
)
run_dirs = list((tmp_path / "work").iterdir())
assert len(run_dirs) == 1
diagnostics = json.loads((run_dirs[0] / "skipped-corrections.json").read_text(encoding="utf-8"))
assert diagnostics["skipped_corrections"][0]["stage"] == "grammar"
assert "more than once" in diagnostics["skipped_corrections"][0]["reason"]