diff --git a/README.md b/README.md index 830b1b5..40346de 100644 --- a/README.md +++ b/README.md @@ -44,6 +44,7 @@ Useful configuration can be supplied by CLI flag or environment variable: - `AUDITA_MAX_SECTION_TOKENS`, default `16000` - `AUDITA_CONFIDENCE_THRESHOLD`, default `0.80` - `AUDITA_MAX_RETRIES`, default `3` +- `AUDITA_GLOSSARY_MAX_LLM_PASSES`, default `3`, for total glossary correction passes - `AUDITA_WORK_DIR`, default `/tmp/audita` `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 2480e38..70e6ae1 100644 --- a/src/audita/chunking.py +++ b/src/audita/chunking.py @@ -71,9 +71,22 @@ def chunk_transcript( if not segments: raise AuditaValidationError("Transcript must contain at least one segment.") - token_estimator = TokenEstimator() if estimator is None else estimator indexed = [IndexedSegment(index=index, segment=segment) for index, segment in enumerate(segments)] + return chunk_indexed_segments(indexed, max_section_tokens, estimator=estimator) + +def chunk_indexed_segments( + indexed_segments: List[IndexedSegment], + max_section_tokens: int, + estimator: Optional[TokenEstimatorProtocol] = None, +) -> List[TranscriptSection]: + if max_section_tokens <= 0: + raise AuditaValidationError("Maximum section token count must be greater than zero.") + if not indexed_segments: + raise AuditaValidationError("Transcript must contain at least one segment.") + + token_estimator = TokenEstimator() if estimator is None else estimator + indexed = list(indexed_segments) sections: List[TranscriptSection] = [] current: List[IndexedSegment] = [] current_tokens = 0 diff --git a/src/audita/cli.py b/src/audita/cli.py index 2961c0d..f78367f 100644 --- a/src/audita/cli.py +++ b/src/audita/cli.py @@ -34,6 +34,7 @@ def _build_parser() -> argparse.ArgumentParser: process.add_argument("--max-section-tokens", type=int, help="maximum estimated tokens per transcript section") process.add_argument("--confidence-threshold", type=float, help="minimum confidence required to apply a correction") process.add_argument("--max-retries", type=int, help="maximum Instructor retries for structured response validation") + process.add_argument("--glossary-max-llm-passes", type=int, help="maximum total LLM passes for glossary corrections") process.add_argument("--work-dir", type=Path, help="directory for per-run scratch diagnostics") return parser @@ -47,6 +48,7 @@ def _process(args: argparse.Namespace) -> int: max_section_tokens=args.max_section_tokens, confidence_threshold=args.confidence_threshold, max_retries=args.max_retries, + glossary_max_llm_passes=args.glossary_max_llm_passes, work_dir=args.work_dir, ) ) diff --git a/src/audita/config.py b/src/audita/config.py index 01068db..e99c78d 100644 --- a/src/audita/config.py +++ b/src/audita/config.py @@ -12,6 +12,7 @@ DEFAULT_MAX_SECTION_TOKENS = 16000 DEFAULT_CONFIDENCE_THRESHOLD = 0.80 DEFAULT_MAX_RETRIES = 3 DEFAULT_WORK_DIR = "/tmp/audita" +DEFAULT_GLOSSARY_MAX_LLM_PASSES = 3 @dataclass(frozen=True) @@ -21,6 +22,7 @@ class ConfigOverrides: max_section_tokens: Optional[int] = None confidence_threshold: Optional[float] = None max_retries: Optional[int] = None + glossary_max_llm_passes: Optional[int] = None work_dir: Optional[Path] = None @@ -32,6 +34,7 @@ class AuditaConfig: max_section_tokens: int = DEFAULT_MAX_SECTION_TOKENS confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD max_retries: int = DEFAULT_MAX_RETRIES + glossary_max_llm_passes: int = DEFAULT_GLOSSARY_MAX_LLM_PASSES work_dir: Path = Path(DEFAULT_WORK_DIR) @classmethod @@ -64,6 +67,12 @@ class AuditaConfig: DEFAULT_MAX_RETRIES, "AUDITA_MAX_RETRIES", ) + glossary_max_llm_passes = _select_int( + selected.glossary_max_llm_passes, + source.get("AUDITA_GLOSSARY_MAX_LLM_PASSES"), + DEFAULT_GLOSSARY_MAX_LLM_PASSES, + "AUDITA_GLOSSARY_MAX_LLM_PASSES", + ) work_dir_value = selected.work_dir or Path(source.get("AUDITA_WORK_DIR") or DEFAULT_WORK_DIR) config = cls( @@ -73,6 +82,7 @@ class AuditaConfig: max_section_tokens=max_section_tokens, confidence_threshold=confidence_threshold, max_retries=max_retries, + glossary_max_llm_passes=glossary_max_llm_passes, work_dir=Path(work_dir_value), ) config.validate() @@ -91,6 +101,8 @@ class AuditaConfig: raise AuditaConfigError("AUDITA_CONFIDENCE_THRESHOLD must be between 0.0 and 1.0.") if self.max_retries < 0: raise AuditaConfigError("AUDITA_MAX_RETRIES must be greater than or equal to zero.") + if self.glossary_max_llm_passes < 1: + raise AuditaConfigError("AUDITA_GLOSSARY_MAX_LLM_PASSES must be greater than or equal to one.") def _get_required_env(env: Mapping[str, str], name: str) -> str: @@ -125,4 +137,3 @@ def _select_float( return float(env_value) except ValueError as exc: raise AuditaConfigError(f"{name} must be a number.") from exc - diff --git a/src/audita/corrections.py b/src/audita/corrections.py index a7e33bd..74a1a81 100644 --- a/src/audita/corrections.py +++ b/src/audita/corrections.py @@ -22,6 +22,8 @@ class SkippedCorrection: class CorrectionApplicationResult: transcript: List[TranscriptSegment] skipped: List[SkippedCorrection] + applied_segment_ids: List[int] + ignored_segment_ids: List[int] def apply_corrections( @@ -34,8 +36,11 @@ def apply_corrections( revised = list(transcript) skipped: List[SkippedCorrection] = [] + applied_segment_ids: List[int] = [] + ignored_segment_ids: List[int] = [] for correction in corrections: if correction.confidence < confidence_threshold: + ignored_segment_ids.append(correction.segment_id) continue reason, actual_text = _target_error(revised, correction) @@ -46,12 +51,13 @@ def apply_corrections( segment = revised[correction.segment_id] revised_text = segment.text.replace(correction.original_text, correction.corrected_text, 1) revised[correction.segment_id] = segment.model_copy(update={"text": revised_text}) + applied_segment_ids.append(correction.segment_id) - indexed_revised = list(enumerate(revised)) - indexed_revised.sort(key=lambda item: (item[1].start, item[1].end, item[0])) return CorrectionApplicationResult( - transcript=[segment for _, segment in indexed_revised], + transcript=revised, skipped=skipped, + applied_segment_ids=applied_segment_ids, + ignored_segment_ids=ignored_segment_ids, ) diff --git a/src/audita/passes.py b/src/audita/passes.py index d2c1c88..97e07b9 100644 --- a/src/audita/passes.py +++ b/src/audita/passes.py @@ -20,6 +20,7 @@ class CorrectionPass(Protocol): glossary: Glossary, config: AuditaConfig, run_dir: Path, + retry_pass: bool = False, ) -> List[CorrectionCandidate]: ... @@ -34,8 +35,9 @@ class GlossaryCorrectionPass: glossary: Glossary, config: AuditaConfig, run_dir: Path, + retry_pass: bool = False, ) -> List[CorrectionCandidate]: - messages = build_glossary_correction_messages(section, glossary) + messages = build_glossary_correction_messages(section, glossary, retry_pass=retry_pass) prompt_path = run_dir / f"prompt-{section.section_index:04d}.json" prompt_path.write_text(json.dumps(messages, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") @@ -43,4 +45,3 @@ class GlossaryCorrectionPass: response_path = run_dir / f"corrections-{section.section_index:04d}.json" response_path.write_text(response.model_dump_json(indent=2) + "\n", encoding="utf-8") return list(response.corrections) - diff --git a/src/audita/pipeline.py b/src/audita/pipeline.py index 3d50457..5865c2b 100644 --- a/src/audita/pipeline.py +++ b/src/audita/pipeline.py @@ -2,10 +2,10 @@ import json import shutil from datetime import datetime from pathlib import Path -from typing import Callable, List, Optional +from typing import Callable, Dict, List, Optional from uuid import uuid4 -from .chunking import TranscriptSection, chunk_transcript +from .chunking import IndexedSegment, TranscriptSection, chunk_indexed_segments from .config import AuditaConfig from .corrections import SkippedCorrection, apply_corrections from .errors import AuditaError @@ -26,44 +26,105 @@ def process_transcript( run_dir = _create_run_dir(config.work_dir) try: _log(progress, f"Created work directory {run_dir}") - sections = chunk_transcript(transcript, config.max_section_tokens) - _write_run_metadata(run_dir, sections, config) + pass_summaries: List[dict] = [] + _write_run_metadata(run_dir, config, pass_summaries) if llm_client is None: from .llm import InstructorLLMClient llm_client = InstructorLLMClient(config) + working = list(transcript) correction_pass = GlossaryCorrectionPass(llm_client) - corrections = [] - for section in sections: - _write_and_validate_section(run_dir, section) - _log( - progress, - f"Processing section {section.section_index + 1}/{len(sections)} " - f"({len(section.segments)} segments, estimated {section.token_count} tokens)", - ) - corrections.extend(correction_pass.run(section, glossary, config, run_dir)) + unresolved_retry_skips: Dict[int, SkippedCorrection] = {} + final_nonretry_skips: List[SkippedCorrection] = [] - application_result = apply_corrections( - transcript, - corrections, - config.confidence_threshold, - ) - _write_skipped_corrections(run_dir, application_result.skipped) - for skipped in application_result.skipped: + for pass_number in range(1, config.glossary_max_llm_passes + 1): + if pass_number == 1: + indexed_segments = _indexed_segments_for_ids(working, list(range(len(working)))) + else: + retry_segment_ids = sorted(unresolved_retry_skips) + if not retry_segment_ids: + break + indexed_segments = _indexed_segments_for_ids(working, retry_segment_ids) + if not indexed_segments: + break + + pass_dir = run_dir / f"pass-{pass_number:04d}" + pass_dir.mkdir() + sections = chunk_indexed_segments(indexed_segments, config.max_section_tokens) + corrections = [] + + for section in sections: + _write_and_validate_section(pass_dir, section) + _log( + progress, + f"Processing glossary pass {pass_number}/{config.glossary_max_llm_passes} " + f"section {section.section_index + 1}/{len(sections)} " + f"({len(section.segments)} segments, estimated {section.token_count} tokens)", + ) + corrections.extend( + correction_pass.run( + section, + glossary, + config, + pass_dir, + retry_pass=pass_number > 1, + ) + ) + + application_result = apply_corrections( + working, + corrections, + config.confidence_threshold, + ) + working = application_result.transcript + + for segment_id in application_result.applied_segment_ids: + unresolved_retry_skips.pop(segment_id, None) + for segment_id in application_result.ignored_segment_ids: + unresolved_retry_skips.pop(segment_id, None) + for skipped in application_result.skipped: + if _is_retryable_skip(skipped, len(working)): + unresolved_retry_skips[skipped.segment_id] = skipped + else: + final_nonretry_skips.append(skipped) + + pass_summaries.append( + { + "pass_number": pass_number, + "retry_pass": pass_number > 1, + "section_count": len(sections), + "segment_count": len(indexed_segments), + "corrections_returned": len(corrections), + "applied_count": len(application_result.applied_segment_ids), + "ignored_below_threshold_count": len(application_result.ignored_segment_ids), + "skipped_count": len(application_result.skipped), + "retry_segment_count": len(unresolved_retry_skips), + } + ) + _write_run_metadata(run_dir, config, pass_summaries) + + if not unresolved_retry_skips: + break + + final_skipped = final_nonretry_skips + [ + unresolved_retry_skips[segment_id] for segment_id in sorted(unresolved_retry_skips) + ] + _write_skipped_corrections(run_dir, final_skipped) + for skipped in final_skipped: _log( progress, f"Skipping correction for segment {skipped.segment_id}: {skipped.reason}", ) - revised = application_result.transcript + revised = _sort_transcript_chronologically(working) 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 - if application_result.skipped: + if final_skipped: _log(progress, f"Skipped correction diagnostics preserved at {run_dir}") else: shutil.rmtree(run_dir) @@ -81,8 +142,8 @@ def _create_run_dir(work_dir: Path) -> Path: def _write_run_metadata( run_dir: Path, - sections: List[TranscriptSection], config: AuditaConfig, + pass_summaries: List[dict], ) -> None: metadata = { "model": config.model, @@ -90,15 +151,8 @@ def _write_run_metadata( "max_section_tokens": config.max_section_tokens, "confidence_threshold": config.confidence_threshold, "max_retries": config.max_retries, - "sections": [ - { - "section_index": section.section_index, - "start_index": section.start_index, - "segment_count": len(section.segments), - "estimated_tokens": section.token_count, - } - for section in sections - ], + "glossary_max_llm_passes": config.glossary_max_llm_passes, + "passes": pass_summaries, } (run_dir / "metadata.json").write_text( json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", @@ -126,6 +180,29 @@ def _write_skipped_corrections(run_dir: Path, skipped: List[SkippedCorrection]) ) +def _indexed_segments_for_ids( + transcript: List[TranscriptSegment], + segment_ids: List[int], +) -> List[IndexedSegment]: + return [ + IndexedSegment(index=segment_id, segment=transcript[segment_id]) + for segment_id in segment_ids + if 0 <= segment_id < len(transcript) + ] + + +def _is_retryable_skip(skipped: SkippedCorrection, transcript_length: int) -> bool: + return 0 <= skipped.segment_id < transcript_length + + +def _sort_transcript_chronologically( + transcript: List[TranscriptSegment], +) -> List[TranscriptSegment]: + indexed = list(enumerate(transcript)) + indexed.sort(key=lambda item: (item[1].start, item[1].end, item[0])) + return [segment for _, segment in indexed] + + 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 6690032..42c8630 100644 --- a/src/audita/prompts.py +++ b/src/audita/prompts.py @@ -8,7 +8,11 @@ from .schemas import Glossary Message = Dict[str, str] -def build_glossary_correction_messages(section: TranscriptSection, glossary: Glossary) -> List[Message]: +def build_glossary_correction_messages( + section: TranscriptSection, + glossary: Glossary, + retry_pass: bool = False, +) -> List[Message]: glossary_json = json.dumps(glossary.model_dump(mode="json"), ensure_ascii=False, indent=2) section_json = json.dumps(section.prompt_payload(), ensure_ascii=False, indent=2) @@ -22,8 +26,17 @@ def build_glossary_correction_messages(section: TranscriptSection, glossary: Glo "Do not rewrite unchanged transcript segments. " "Preserve speaker names, timestamps, and meaning." ) + retry_guidance = "" + if retry_pass: + retry_guidance = ( + "Retry guidance:\n" + "These segments are being retried because previous correction spans did not apply cleanly. " + "Copy original_text exactly from the current segment text, using only the span that needs replacement.\n\n" + ) + user = ( "Review this transcript section and return only corrections that should be applied.\n\n" + f"{retry_guidance}" "Rules:\n" "- Correct domain-specific names, aliases, jargon, deities, locations, NPCs, players, and similar terms only when both the glossary and surrounding transcript context support the correction.\n" "- The correction must be likely to fix a transcription error: the original words should be phonetically or acoustically similar to the corrected words in spoken English.\n" diff --git a/tests/test_cli.py b/tests/test_cli.py index fa5daab..e237716 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -10,3 +10,10 @@ def test_cli_help_uses_audita_program_name(capsys): assert exc.value.code == 0 assert capsys.readouterr().out.startswith("usage: audita ") + +def test_process_help_includes_glossary_pass_flag(capsys): + with pytest.raises(SystemExit) as exc: + main(["process", "--help"]) + + assert exc.value.code == 0 + assert "--glossary-max-llm-passes" in capsys.readouterr().out diff --git a/tests/test_config.py b/tests/test_config.py index 28f6137..1bb9166 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -3,7 +3,7 @@ from pathlib import Path import pytest from audita.config import AuditaConfig, ConfigOverrides -from audita.config import DEFAULT_MAX_RETRIES, DEFAULT_MAX_SECTION_TOKENS, DEFAULT_WORK_DIR +from audita.config import DEFAULT_GLOSSARY_MAX_LLM_PASSES, DEFAULT_MAX_RETRIES, DEFAULT_MAX_SECTION_TOKENS, DEFAULT_WORK_DIR from audita.errors import AuditaConfigError @@ -12,6 +12,7 @@ def test_config_uses_defaults_with_api_key(): assert config.max_section_tokens == DEFAULT_MAX_SECTION_TOKENS assert config.max_retries == DEFAULT_MAX_RETRIES + assert config.glossary_max_llm_passes == DEFAULT_GLOSSARY_MAX_LLM_PASSES assert config.work_dir == Path(DEFAULT_WORK_DIR) @@ -22,6 +23,7 @@ def test_config_env_overrides_defaults(): "AUDITA_MAX_SECTION_TOKENS": "42", "AUDITA_CONFIDENCE_THRESHOLD": "0.9", "AUDITA_MAX_RETRIES": "5", + "AUDITA_GLOSSARY_MAX_LLM_PASSES": "7", "AUDITA_WORK_DIR": "/tmp/custom-audita", } ) @@ -29,6 +31,7 @@ def test_config_env_overrides_defaults(): assert config.max_section_tokens == 42 assert config.confidence_threshold == 0.9 assert config.max_retries == 5 + assert config.glossary_max_llm_passes == 7 assert config.work_dir == Path("/tmp/custom-audita") @@ -38,17 +41,20 @@ def test_config_cli_overrides_env(): "OPENROUTER_API_KEY": "key", "AUDITA_MAX_SECTION_TOKENS": "42", "AUDITA_MAX_RETRIES": "5", + "AUDITA_GLOSSARY_MAX_LLM_PASSES": "7", "AUDITA_WORK_DIR": "/tmp/env-audita", }, overrides=ConfigOverrides( max_section_tokens=100, max_retries=3, + glossary_max_llm_passes=2, work_dir=Path("/tmp/cli-audita"), ), ) assert config.max_section_tokens == 100 assert config.max_retries == 3 + assert config.glossary_max_llm_passes == 2 assert config.work_dir == Path("/tmp/cli-audita") @@ -63,3 +69,9 @@ def test_config_rejects_bad_env_int(): env={"OPENROUTER_API_KEY": "key", "AUDITA_MAX_SECTION_TOKENS": "many"} ) + +def test_config_rejects_invalid_glossary_pass_count(): + with pytest.raises(AuditaConfigError): + AuditaConfig.from_sources( + env={"OPENROUTER_API_KEY": "key", "AUDITA_GLOSSARY_MAX_LLM_PASSES": "0"} + ) diff --git a/tests/test_corrections.py b/tests/test_corrections.py index c4b663a..edd0c98 100644 --- a/tests/test_corrections.py +++ b/tests/test_corrections.py @@ -16,7 +16,7 @@ def _transcript(): ) -def test_apply_corrections_uses_threshold_and_sorts_chronologically(): +def test_apply_corrections_uses_threshold_and_preserves_segment_id_order(): transcript = _transcript() corrections = [ CorrectionCandidate( @@ -29,8 +29,8 @@ def test_apply_corrections_uses_threshold_and_sorts_chronologically(): result = apply_corrections(transcript, corrections, confidence_threshold=0.8) - assert [segment.speaker for segment in result.transcript] == ["Mike", "Eric"] - assert result.transcript[1].text == "I ask Chauntea for help." + assert [segment.speaker for segment in result.transcript] == ["Eric", "Mike"] + assert result.transcript[0].text == "I ask Chauntea for help." assert result.skipped == [] @@ -47,7 +47,7 @@ def test_apply_corrections_ignores_below_threshold(): result = apply_corrections(transcript, corrections, confidence_threshold=0.8) - assert result.transcript[1].text == "I ask Chontia for help." + assert result.transcript[0].text == "I ask Chontia for help." assert result.skipped == [] @@ -68,7 +68,7 @@ def test_apply_corrections_allows_multiple_distinct_spans_in_one_segment(): result = apply_corrections(transcript, [first, second], confidence_threshold=0.8) - assert result.transcript[1].text == "I ask Chauntea for guidance." + assert result.transcript[0].text == "I ask Chauntea for guidance." assert result.skipped == [] @@ -83,7 +83,7 @@ def test_apply_corrections_skips_missing_substring(): result = apply_corrections(transcript, [correction], confidence_threshold=0.8) - assert result.transcript[1].text == "I ask Chontia for help." + assert result.transcript[0].text == "I ask Chontia for help." assert len(result.skipped) == 1 assert result.skipped[0].segment_id == 0 assert result.skipped[0].actual_text == "I ask Chontia for help." @@ -101,7 +101,7 @@ def test_apply_corrections_skips_missing_segment_id(): result = apply_corrections(transcript, [correction], confidence_threshold=0.8) - assert [segment.text for segment in result.transcript] == ["Then Lyra.", "I ask Chontia for help."] + assert [segment.text for segment in result.transcript] == ["I ask Chontia for help.", "Then Lyra."] assert len(result.skipped) == 1 assert result.skipped[0].segment_id == 99 assert "does not exist" in result.skipped[0].reason @@ -118,7 +118,7 @@ def test_apply_corrections_skips_no_op(): result = apply_corrections(transcript, [correction], confidence_threshold=0.8) - assert result.transcript[1].text == "I ask Chontia for help." + assert result.transcript[0].text == "I ask Chontia for help." assert len(result.skipped) == 1 assert "identical" in result.skipped[0].reason @@ -156,7 +156,7 @@ def test_apply_corrections_skips_empty_original_text(): result = apply_corrections(transcript, [correction], confidence_threshold=0.8) - assert result.transcript[1].text == "I ask Chontia for help." + assert result.transcript[0].text == "I ask Chontia for help." assert len(result.skipped) == 1 assert "empty" in result.skipped[0].reason diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index e1c946e..4986aa6 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -9,18 +9,21 @@ 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): +def _config(tmp_path, glossary_max_llm_passes=3): return AuditaConfig( api_key="key", max_section_tokens=16000, confidence_threshold=0.8, max_retries=3, + glossary_max_llm_passes=glossary_max_llm_passes, work_dir=tmp_path / "work", ) @@ -40,7 +43,8 @@ def _transcript(): return parse_transcript_json( """ [ - {"speaker": "Eric", "start": 0.0, "end": 1.0, "text": "I ask Chontia."} + {"speaker": "Eric", "start": 10.0, "end": 11.0, "text": "I ask Chontia."}, + {"speaker": "Mike", "start": 0.0, "end": 1.0, "text": "Then Lyra."} ] """ ) @@ -62,7 +66,8 @@ def test_pipeline_processes_with_fake_llm_and_cleans_work_dir(tmp_path): llm_client=fake_client, ) - assert revised[0].text == "I ask Chauntea." + assert [segment.speaker for segment in revised] == ["Mike", "Eric"] + assert revised[1].text == "I ask Chauntea." assert fake_client.calls == 1 assert list((tmp_path / "work").iterdir()) == [] @@ -80,12 +85,12 @@ def test_pipeline_skips_bad_correction_and_preserves_diagnostics(tmp_path): revised = process_transcript( _transcript(), _glossary(), - _config(tmp_path), + _config(tmp_path, glossary_max_llm_passes=1), llm_client=fake_client, progress=progress.append, ) - assert revised[0].text == "I ask Chontia." + assert revised[1].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 @@ -94,3 +99,106 @@ def test_pipeline_skips_bad_correction_and_preserves_diagnostics(tmp_path): diagnostics = json.loads(skipped_path.read_text(encoding="utf-8")) assert diagnostics["skipped_corrections"][0]["segment_id"] == 0 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( + segment_id=0, + original_text="Contia", + corrected_text="Chauntea", + confidence=0.95, + ) + second_pass = CorrectionCandidate( + segment_id=0, + original_text="Chontia", + corrected_text="Chauntea", + confidence=0.95, + ) + fake_client = FakeLLMClient( + [ + CorrectionSet(corrections=[first_pass]), + CorrectionSet(corrections=[second_pass]), + ] + ) + + revised = process_transcript( + _transcript(), + _glossary(), + _config(tmp_path, glossary_max_llm_passes=3), + llm_client=fake_client, + ) + + assert fake_client.calls == 2 + assert [segment.speaker for segment in revised] == ["Mike", "Eric"] + assert revised[1].text == "I ask Chauntea." + assert list((tmp_path / "work").iterdir()) == [] + + +def test_pipeline_retry_prompt_contains_only_valid_deduped_segment_ids(tmp_path): + first_bad = CorrectionCandidate( + segment_id=0, + original_text="Contia", + corrected_text="Chauntea", + confidence=0.95, + ) + second_bad_same_segment = CorrectionCandidate( + segment_id=0, + original_text="Still wrong", + corrected_text="Chauntea", + confidence=0.95, + ) + invalid_segment = CorrectionCandidate( + segment_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=[]), + ] + ) + + process_transcript( + _transcript(), + _glossary(), + _config(tmp_path, glossary_max_llm_passes=2), + llm_client=fake_client, + ) + + assert fake_client.calls == 2 + retry_prompt = fake_client.messages[1][1]["content"] + retry_payload = json.loads(retry_prompt.split("Transcript section:\n", maxsplit=1)[1]) + assert retry_payload == [{"segment_id": 0, "original_text": "I ask Chontia."}] + assert "Retry guidance" in retry_prompt + + +def test_pipeline_writes_pass_metadata_for_unresolved_retries(tmp_path): + first_pass = CorrectionCandidate( + segment_id=0, + original_text="Contia", + corrected_text="Chauntea", + confidence=0.95, + ) + fake_client = FakeLLMClient( + [ + CorrectionSet(corrections=[first_pass]), + 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 + 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