audita/tests/test_pipeline.py

import json

from audita.config import AuditaConfig
from audita.pipeline import process_transcript
from audita.schemas import CorrectionCandidate, CorrectionSet, parse_glossary_yaml, parse_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):
    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",
    )


def _glossary():
    return parse_glossary_yaml(
        """
        glossary:
          - name: "Chauntea"
            category: deity
            summary: "Chauntea is a deity."
        """
    )


def _transcript():
    return 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."}
        ]
        """
    )


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])])

    revised = process_transcript(
        _transcript(),
        _glossary(),
        _config(tmp_path),
        llm_client=fake_client,
    )

    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()) == []


def test_pipeline_skips_bad_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])])
    progress = []

    revised = process_transcript(
        _transcript(),
        _glossary(),
        _config(tmp_path, glossary_max_llm_passes=1),
        llm_client=fake_client,
        progress=progress.append,
    )

    assert revised[1].text == "I ask Chontia."
    assert any("Skipping 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]["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]),
        ]
    )

    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_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=[]),
        ]
    )

    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 == [{"id": 1, "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(
        id=1,
        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