Implemented a second LLM stage for grammatical review
This commit is contained in:
@@ -17,13 +17,15 @@ class FakeLLMClient:
|
||||
return self.responses.pop(0)
|
||||
|
||||
|
||||
def _config(tmp_path, glossary_max_llm_passes=3):
|
||||
def _config(tmp_path, glossary_max_llm_passes=3, grammar_max_llm_passes=3):
|
||||
return AuditaConfig(
|
||||
api_key="key",
|
||||
max_section_tokens=16000,
|
||||
confidence_threshold=0.8,
|
||||
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",
|
||||
)
|
||||
|
||||
@@ -57,7 +59,12 @@ def test_pipeline_processes_with_fake_llm_and_cleans_work_dir(tmp_path):
|
||||
corrected_text="Chauntea",
|
||||
confidence=0.95,
|
||||
)
|
||||
fake_client = FakeLLMClient([CorrectionSet(corrections=[correction])])
|
||||
fake_client = FakeLLMClient(
|
||||
[
|
||||
CorrectionSet(corrections=[correction]),
|
||||
CorrectionSet(corrections=[]),
|
||||
]
|
||||
)
|
||||
|
||||
revised = process_transcript(
|
||||
_transcript(),
|
||||
@@ -68,18 +75,23 @@ def test_pipeline_processes_with_fake_llm_and_cleans_work_dir(tmp_path):
|
||||
|
||||
assert [segment.speaker for segment in revised] == ["Mike", "Eric"]
|
||||
assert revised[1].text == "I ask Chauntea."
|
||||
assert fake_client.calls == 1
|
||||
assert fake_client.calls == 2
|
||||
assert list((tmp_path / "work").iterdir()) == []
|
||||
|
||||
|
||||
def test_pipeline_skips_bad_correction_and_preserves_diagnostics(tmp_path):
|
||||
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])])
|
||||
fake_client = FakeLLMClient(
|
||||
[
|
||||
CorrectionSet(corrections=[correction]),
|
||||
CorrectionSet(corrections=[]),
|
||||
]
|
||||
)
|
||||
progress = []
|
||||
|
||||
revised = process_transcript(
|
||||
@@ -91,12 +103,13 @@ def test_pipeline_skips_bad_correction_and_preserves_diagnostics(tmp_path):
|
||||
)
|
||||
|
||||
assert revised[1].text == "I ask Chontia."
|
||||
assert any("Skipping correction for id 1" in message for message in progress)
|
||||
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"]
|
||||
|
||||
@@ -118,6 +131,7 @@ def test_pipeline_retries_skipped_segment_and_cleans_work_dir_when_fixed(tmp_pat
|
||||
[
|
||||
CorrectionSet(corrections=[first_pass]),
|
||||
CorrectionSet(corrections=[second_pass]),
|
||||
CorrectionSet(corrections=[]),
|
||||
]
|
||||
)
|
||||
|
||||
@@ -128,7 +142,7 @@ def test_pipeline_retries_skipped_segment_and_cleans_work_dir_when_fixed(tmp_pat
|
||||
llm_client=fake_client,
|
||||
)
|
||||
|
||||
assert fake_client.calls == 2
|
||||
assert fake_client.calls == 3
|
||||
assert [segment.speaker for segment in revised] == ["Mike", "Eric"]
|
||||
assert revised[1].text == "I ask Chauntea."
|
||||
assert list((tmp_path / "work").iterdir()) == []
|
||||
@@ -157,6 +171,7 @@ def test_pipeline_retry_prompt_contains_only_valid_deduped_ids(tmp_path):
|
||||
[
|
||||
CorrectionSet(corrections=[first_bad, second_bad_same_segment, invalid_segment]),
|
||||
CorrectionSet(corrections=[]),
|
||||
CorrectionSet(corrections=[]),
|
||||
]
|
||||
)
|
||||
|
||||
@@ -167,14 +182,14 @@ def test_pipeline_retry_prompt_contains_only_valid_deduped_ids(tmp_path):
|
||||
llm_client=fake_client,
|
||||
)
|
||||
|
||||
assert fake_client.calls == 2
|
||||
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_pass_metadata_for_unresolved_retries(tmp_path):
|
||||
def test_pipeline_writes_stage_metadata_for_unresolved_retries(tmp_path):
|
||||
first_pass = CorrectionCandidate(
|
||||
id=1,
|
||||
original_text="Contia",
|
||||
@@ -185,6 +200,7 @@ def test_pipeline_writes_pass_metadata_for_unresolved_retries(tmp_path):
|
||||
[
|
||||
CorrectionSet(corrections=[first_pass]),
|
||||
CorrectionSet(corrections=[]),
|
||||
CorrectionSet(corrections=[]),
|
||||
]
|
||||
)
|
||||
|
||||
@@ -197,8 +213,150 @@ def test_pipeline_writes_pass_metadata_for_unresolved_retries(tmp_path):
|
||||
|
||||
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()
|
||||
metadata = json.loads((run_dirs[0] / "metadata.json").read_text(encoding="utf-8"))
|
||||
assert metadata["glossary_max_llm_passes"] == 2
|
||||
assert [item["pass_number"] for item in metadata["passes"]] == [1, 2]
|
||||
assert metadata["passes"][0]["retry_segment_count"] == 1
|
||||
assert metadata["passes"][1]["retry_pass"] is True
|
||||
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_transcript_json(
|
||||
"""
|
||||
[
|
||||
{"id": 1, "speaker": "Eric", "start": 10.0, "end": 11.0, "text": "i ask Chontia."},
|
||||
{"id": 2, "speaker": "Mike", "start": 0.0, "end": 1.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[1].text == "I ask Chauntea."
|
||||
|
||||
|
||||
def test_grammar_stage_retries_repeated_span_and_applies_unique_retry(tmp_path):
|
||||
transcript = parse_transcript_json(
|
||||
"""
|
||||
[
|
||||
{"id": 1, "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[1].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"]
|
||||
|
||||
Reference in New Issue
Block a user