From 23532cade1db177beb60dd5a64320dc91a6035f2 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Tue, 21 Apr 2026 11:14:56 -0500 Subject: [PATCH] Simplified the json schema passed to the LLM, and implemented more forgiving error handling for LLM proposed corrections --- README.md | 2 +- src/audita/chunking.py | 9 ++-- src/audita/corrections.py | 96 ++++++++++++++++++++++++++++----------- src/audita/pipeline.py | 36 +++++++++++++-- src/audita/prompts.py | 5 +- src/audita/schemas.py | 12 ++--- tests/test_corrections.py | 82 +++++++++++++++++++++------------ tests/test_pipeline.py | 41 ++++++++--------- tests/test_prompts.py | 30 ++++++++++++ 9 files changed, 216 insertions(+), 97 deletions(-) diff --git a/README.md b/README.md index 0fd7b68..830b1b5 100644 --- a/README.md +++ b/README.md @@ -46,4 +46,4 @@ Useful configuration can be supplied by CLI flag or environment variable: - `AUDITA_MAX_RETRIES`, default `3` - `AUDITA_WORK_DIR`, default `/tmp/audita` -`AUDITA_WORK_DIR` stores per-run diagnostics while processing. Successful runs clean up their run directory; failed runs preserve it for debugging. +`AUDITA_WORK_DIR` stores per-run diagnostics while processing. Successful runs clean up their run directory unless corrections are skipped for target mismatches; failed runs and skipped-correction runs preserve diagnostics for debugging. diff --git a/src/audita/chunking.py b/src/audita/chunking.py index 821c133..2480e38 100644 --- a/src/audita/chunking.py +++ b/src/audita/chunking.py @@ -41,9 +41,7 @@ class IndexedSegment: return self.segment.model_dump(mode="json") def prompt_payload(self) -> dict: - payload = self.transcript_payload() - payload["segment_index"] = self.index - return payload + return {"segment_id": self.index, "original_text": self.segment.text} @dataclass(frozen=True) @@ -81,7 +79,7 @@ def chunk_transcript( current_tokens = 0 for item in indexed: - single_payload = [item.transcript_payload()] + single_payload = [item.prompt_payload()] single_tokens = token_estimator.estimate_json(single_payload) if single_tokens > max_section_tokens: raise AuditaValidationError( @@ -91,7 +89,7 @@ def chunk_transcript( candidate = current + [item] candidate_tokens = token_estimator.estimate_json( - [candidate_item.transcript_payload() for candidate_item in candidate] + [candidate_item.prompt_payload() for candidate_item in candidate] ) if current and candidate_tokens > max_section_tokens: sections.append(_make_section(len(sections), current, current_tokens)) @@ -121,4 +119,3 @@ def _make_section( segments=list(segments), token_count=token_count, ) - diff --git a/src/audita/corrections.py b/src/audita/corrections.py index adcb225..3ac4540 100644 --- a/src/audita/corrections.py +++ b/src/audita/corrections.py @@ -1,29 +1,60 @@ -from typing import Dict, Iterable, List +from dataclasses import asdict, dataclass +from typing import Dict, Iterable, List, Optional, Tuple from .errors import AuditaValidationError from .schemas import CorrectionCandidate, TranscriptSegment +@dataclass(frozen=True) +class SkippedCorrection: + segment_id: int + reason: str + original_text: str + corrected_text: str + confidence: float + actual_text: Optional[str] = None + + def to_dict(self) -> dict: + return asdict(self) + + +@dataclass(frozen=True) +class CorrectionApplicationResult: + transcript: List[TranscriptSegment] + skipped: List[SkippedCorrection] + + def apply_corrections( transcript: List[TranscriptSegment], corrections: Iterable[CorrectionCandidate], confidence_threshold: float, -) -> List[TranscriptSegment]: +) -> CorrectionApplicationResult: if not 0.0 <= confidence_threshold <= 1.0: raise AuditaValidationError("Confidence threshold must be between 0.0 and 1.0.") - correction_by_index: Dict[int, CorrectionCandidate] = {} + correction_by_id: Dict[int, CorrectionCandidate] = {} + skipped: List[SkippedCorrection] = [] for correction in corrections: - if correction.segment_index in correction_by_index: - raise AuditaValidationError( - f"Duplicate corrections returned for segment {correction.segment_index}." + segment_id = correction.segment_id + if segment_id in correction_by_id: + skipped.append( + _skip( + correction, + "duplicate correction for segment already handled", + ) ) - _validate_target(transcript, correction) - correction_by_index[correction.segment_index] = correction + continue + + reason, actual_text = _target_error(transcript, correction) + if reason is not None: + skipped.append(_skip(correction, reason, actual_text=actual_text)) + continue + + correction_by_id[segment_id] = correction revised: List[TranscriptSegment] = [] - for original_index, segment in enumerate(transcript): - correction = correction_by_index.get(original_index) + for original_id, segment in enumerate(transcript): + correction = correction_by_id.get(original_id) if correction is not None and correction.confidence >= confidence_threshold: revised.append(segment.model_copy(update={"text": correction.corrected_text})) else: @@ -31,23 +62,36 @@ def apply_corrections( indexed_revised = list(enumerate(revised)) indexed_revised.sort(key=lambda item: (item[1].start, item[1].end, item[0])) - return [segment for _, segment in indexed_revised] + return CorrectionApplicationResult( + transcript=[segment for _, segment in indexed_revised], + skipped=skipped, + ) -def _validate_target(transcript: List[TranscriptSegment], correction: CorrectionCandidate) -> None: - if correction.segment_index >= len(transcript): - raise AuditaValidationError( - f"Correction targets missing segment {correction.segment_index}." - ) +def _target_error( + transcript: List[TranscriptSegment], + correction: CorrectionCandidate, +) -> Tuple[Optional[str], Optional[str]]: + if correction.segment_id >= len(transcript): + return "segment_id does not exist in transcript", None - segment = transcript[correction.segment_index] - if ( - correction.speaker != segment.speaker - or correction.start != segment.start - or correction.end != segment.end - or correction.original_text != segment.text - ): - raise AuditaValidationError( - f"Correction target does not exactly match segment {correction.segment_index}." - ) + segment = transcript[correction.segment_id] + if correction.original_text != segment.text: + return "original_text does not exactly match segment text", segment.text + return None, None + + +def _skip( + correction: CorrectionCandidate, + reason: str, + actual_text: Optional[str] = None, +) -> SkippedCorrection: + return SkippedCorrection( + segment_id=correction.segment_id, + reason=reason, + original_text=correction.original_text, + corrected_text=correction.corrected_text, + confidence=correction.confidence, + actual_text=actual_text, + ) diff --git a/src/audita/pipeline.py b/src/audita/pipeline.py index f3ef352..3d50457 100644 --- a/src/audita/pipeline.py +++ b/src/audita/pipeline.py @@ -7,7 +7,7 @@ from uuid import uuid4 from .chunking import TranscriptSection, chunk_transcript from .config import AuditaConfig -from .corrections import apply_corrections +from .corrections import SkippedCorrection, apply_corrections from .errors import AuditaError from .passes import GlossaryCorrectionPass, LLMClient from .schemas import Glossary, TranscriptSegment, parse_transcript_json @@ -45,15 +45,29 @@ def process_transcript( ) corrections.extend(correction_pass.run(section, glossary, config, run_dir)) - revised = apply_corrections(transcript, corrections, config.confidence_threshold) + application_result = apply_corrections( + transcript, + corrections, + config.confidence_threshold, + ) + _write_skipped_corrections(run_dir, application_result.skipped) + for skipped in application_result.skipped: + _log( + progress, + f"Skipping correction for segment {skipped.segment_id}: {skipped.reason}", + ) + revised = application_result.transcript except Exception as exc: message = f"{exc} Diagnostics preserved at {run_dir}" if isinstance(exc, AuditaError): raise type(exc)(message) from exc raise AuditaError(message) from exc - shutil.rmtree(run_dir) - _log(progress, "Removed work directory after successful run") + if application_result.skipped: + _log(progress, f"Skipped correction diagnostics preserved at {run_dir}") + else: + shutil.rmtree(run_dir) + _log(progress, "Removed work directory after successful run") return revised @@ -99,7 +113,19 @@ def _write_and_validate_section(run_dir: Path, section: TranscriptSection) -> No parse_transcript_json(section_json) +def _write_skipped_corrections(run_dir: Path, skipped: List[SkippedCorrection]) -> None: + skipped_path = run_dir / "skipped-corrections.json" + skipped_path.write_text( + json.dumps( + {"skipped_corrections": [item.to_dict() for item in skipped]}, + ensure_ascii=False, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + + def _log(progress: Optional[ProgressCallback], message: str) -> None: if progress is not None: progress(message) - diff --git a/src/audita/prompts.py b/src/audita/prompts.py index 7e3c143..20bde2a 100644 --- a/src/audita/prompts.py +++ b/src/audita/prompts.py @@ -31,7 +31,10 @@ def build_glossary_correction_messages(section: TranscriptSection, glossary: Glo "- Inappropriate example: correcting \"Lyra\" to \"Jesters\" should be omitted because those words are not similar in spoken English, even if \"Jesters\" appears in the glossary.\n" "- Do not replace one clear glossary term, character name, location, or ordinary word with a different glossary term unless it is a plausible mishearing.\n" "- Assign high confidence only when the correction is supported by glossary evidence, local context, and spoken-word similarity; otherwise omit the correction.\n" - "- Use the exact segment_index, speaker, start, end, and original_text from the input segment.\n" + "- Use the exact segment_id and original_text from the input segment.\n" + "- Each returned correction must contain only segment_id, original_text, corrected_text, and confidence.\n" + "- Do not return speaker, start, or end fields.\n" + "- Return only changed segments; do not return entries for unchanged segments.\n" "- corrected_text must contain the full corrected text for that segment.\n" "- confidence must be between 0.0 and 1.0.\n" "- If no corrections are needed, return an empty corrections list.\n\n" diff --git a/src/audita/schemas.py b/src/audita/schemas.py index d88ff5d..71f460e 100644 --- a/src/audita/schemas.py +++ b/src/audita/schemas.py @@ -82,22 +82,19 @@ class Glossary(BaseModel): class CorrectionCandidate(BaseModel): model_config = ConfigDict(extra="forbid") - segment_index: int = Field(ge=0) - speaker: StrictStr - start: float - end: float + segment_id: int = Field(ge=0) original_text: StrictStr corrected_text: StrictStr confidence: float = Field(ge=0.0, le=1.0) - @field_validator("segment_index", mode="before") + @field_validator("segment_id", mode="before") @classmethod - def require_integer_index(cls, value: Any) -> int: + def require_integer_id(cls, value: Any) -> int: if isinstance(value, bool) or not isinstance(value, int): raise ValueError("must be an integer") return value - @field_validator("start", "end", "confidence", mode="before") + @field_validator("confidence", mode="before") @classmethod def require_number(cls, value: Any) -> float: if isinstance(value, bool) or not isinstance(value, (int, float)): @@ -158,4 +155,3 @@ def parse_glossary_yaml(raw: str) -> Glossary: def transcript_to_json(segments: List[TranscriptSegment]) -> str: payload = [segment.model_dump(mode="json") for segment in segments] return json.dumps(payload, ensure_ascii=False, indent=2) + "\n" - diff --git a/tests/test_corrections.py b/tests/test_corrections.py index 97e0e41..5f5bd73 100644 --- a/tests/test_corrections.py +++ b/tests/test_corrections.py @@ -20,69 +20,95 @@ def test_apply_corrections_uses_threshold_and_sorts_chronologically(): transcript = _transcript() corrections = [ CorrectionCandidate( - segment_index=0, - speaker="Eric", - start=10.0, - end=11.0, + segment_id=0, original_text="I ask Chontia.", corrected_text="I ask Chauntea.", confidence=0.8, ) ] - revised = apply_corrections(transcript, corrections, confidence_threshold=0.8) + result = apply_corrections(transcript, corrections, confidence_threshold=0.8) - assert [segment.speaker for segment in revised] == ["Mike", "Eric"] - assert revised[1].text == "I ask Chauntea." + assert [segment.speaker for segment in result.transcript] == ["Mike", "Eric"] + assert result.transcript[1].text == "I ask Chauntea." + assert result.skipped == [] def test_apply_corrections_ignores_below_threshold(): transcript = _transcript() corrections = [ CorrectionCandidate( - segment_index=0, - speaker="Eric", - start=10.0, - end=11.0, + segment_id=0, original_text="I ask Chontia.", corrected_text="I ask Chauntea.", confidence=0.79, ) ] - revised = apply_corrections(transcript, corrections, confidence_threshold=0.8) + result = apply_corrections(transcript, corrections, confidence_threshold=0.8) - assert revised[1].text == "I ask Chontia." + assert result.transcript[1].text == "I ask Chontia." + assert result.skipped == [] -def test_apply_corrections_rejects_duplicate_targets(): +def test_apply_corrections_skips_duplicate_targets(): transcript = _transcript() - correction = CorrectionCandidate( - segment_index=0, - speaker="Eric", - start=10.0, - end=11.0, + first = CorrectionCandidate( + segment_id=0, original_text="I ask Chontia.", corrected_text="I ask Chauntea.", confidence=0.8, ) + second = CorrectionCandidate( + segment_id=0, + original_text="I ask Chontia.", + corrected_text="I ask Something Else.", + confidence=0.9, + ) - with pytest.raises(AuditaValidationError): - apply_corrections(transcript, [correction, correction], confidence_threshold=0.8) + result = apply_corrections(transcript, [first, second], confidence_threshold=0.8) + + assert result.transcript[1].text == "I ask Chauntea." + assert len(result.skipped) == 1 + assert result.skipped[0].segment_id == 0 + assert "duplicate" in result.skipped[0].reason -def test_apply_corrections_rejects_mismatched_original_text(): +def test_apply_corrections_skips_mismatched_original_text(): transcript = _transcript() correction = CorrectionCandidate( - segment_index=0, - speaker="Eric", - start=10.0, - end=11.0, + segment_id=0, original_text="Different text.", corrected_text="I ask Chauntea.", confidence=0.8, ) - with pytest.raises(AuditaValidationError): - apply_corrections(transcript, [correction], confidence_threshold=0.8) + result = apply_corrections(transcript, [correction], confidence_threshold=0.8) + assert result.transcript[1].text == "I ask Chontia." + assert len(result.skipped) == 1 + assert result.skipped[0].segment_id == 0 + assert result.skipped[0].actual_text == "I ask Chontia." + assert "original_text" in result.skipped[0].reason + + +def test_apply_corrections_skips_missing_segment_id(): + transcript = _transcript() + correction = CorrectionCandidate( + segment_id=99, + original_text="Missing.", + corrected_text="Still missing.", + confidence=0.8, + ) + + result = apply_corrections(transcript, [correction], confidence_threshold=0.8) + + assert [segment.text for segment in result.transcript] == ["Then Lyra.", "I ask Chontia."] + assert len(result.skipped) == 1 + assert result.skipped[0].segment_id == 99 + assert "does not exist" in result.skipped[0].reason + + +def test_apply_corrections_rejects_invalid_threshold(): + with pytest.raises(AuditaValidationError): + apply_corrections(_transcript(), [], confidence_threshold=1.1) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 1ab3ecb..9effce4 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,9 +1,6 @@ -from pathlib import Path - -import pytest +import json from audita.config import AuditaConfig -from audita.errors import AuditaValidationError from audita.pipeline import process_transcript from audita.schemas import CorrectionCandidate, CorrectionSet, parse_glossary_yaml, parse_transcript_json @@ -51,10 +48,7 @@ def _transcript(): def test_pipeline_processes_with_fake_llm_and_cleans_work_dir(tmp_path): correction = CorrectionCandidate( - segment_index=0, - speaker="Eric", - start=0.0, - end=1.0, + segment_id=0, original_text="I ask Chontia.", corrected_text="I ask Chauntea.", confidence=0.95, @@ -73,27 +67,30 @@ def test_pipeline_processes_with_fake_llm_and_cleans_work_dir(tmp_path): assert list((tmp_path / "work").iterdir()) == [] -def test_pipeline_preserves_work_dir_on_failure(tmp_path): +def test_pipeline_skips_bad_correction_and_preserves_diagnostics(tmp_path): correction = CorrectionCandidate( - segment_index=0, - speaker="Eric", - start=0.0, - end=1.0, + segment_id=0, original_text="Different text.", corrected_text="I ask Chauntea.", confidence=0.95, ) fake_client = FakeLLMClient([CorrectionSet(corrections=[correction])]) + progress = [] - with pytest.raises(AuditaValidationError): - process_transcript( - _transcript(), - _glossary(), - _config(tmp_path), - llm_client=fake_client, - ) + revised = process_transcript( + _transcript(), + _glossary(), + _config(tmp_path), + llm_client=fake_client, + progress=progress.append, + ) + assert revised[0].text == "I ask Chontia." + assert any("Skipping correction for segment 0" in message for message in progress) preserved = list((tmp_path / "work").iterdir()) assert len(preserved) == 1 - assert (Path(preserved[0]) / "section-0000.json").exists() - + 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]["segment_id"] == 0 + assert "original_text" in diagnostics["skipped_corrections"][0]["reason"] diff --git a/tests/test_prompts.py b/tests/test_prompts.py index b9b77f7..902a15f 100644 --- a/tests/test_prompts.py +++ b/tests/test_prompts.py @@ -1,3 +1,5 @@ +import json + from audita.chunking import chunk_transcript from audita.prompts import build_glossary_correction_messages from audita.schemas import parse_glossary_yaml, parse_transcript_json @@ -32,3 +34,31 @@ def test_prompt_requires_acoustically_plausible_transcription_errors(): assert '"gestures" to "Jesters"' in prompt_text assert '"Lyra" to "Jesters"' in prompt_text assert "should be omitted" in prompt_text + + +def test_prompt_uses_simplified_segment_payload(): + transcript = parse_transcript_json( + """ + [ + {"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "The gestures are nearby."} + ] + """ + ) + glossary = parse_glossary_yaml( + """ + glossary: + - name: "Jesters" + category: faction + summary: "The Jesters are a local faction." + """ + ) + section = chunk_transcript(transcript, max_section_tokens=16000)[0] + + messages = build_glossary_correction_messages(section, glossary) + transcript_json = messages[1]["content"].split("Transcript section:\n", maxsplit=1)[1] + prompt_segments = json.loads(transcript_json) + + assert prompt_segments == [{"segment_id": 0, "original_text": "The gestures are nearby."}] + assert "speaker" not in prompt_segments[0] + assert "start" not in prompt_segments[0] + assert "end" not in prompt_segments[0]