643 lines
20 KiB
Python
643 lines
20 KiB
Python
import threading
|
|
|
|
from audita.core.chunking import IndexedSegment, TranscriptSection
|
|
from audita.core.config import AuditaConfig, ConfigOverrides
|
|
from audita.core.errors import AuditaLLMError
|
|
from audita.core.schemas import parse_glossary_yaml, parse_transcript_json
|
|
from audita.framework.models import CorrectionProposal, ModuleContext, ModuleRunSpec
|
|
from audita.framework.runner import PipelineRunner
|
|
from audita.validators import (
|
|
MeaningReversalValidator,
|
|
ProposalConfidenceValidator,
|
|
ProtectedGlossaryTermsValidator,
|
|
SpokenFormPlausibilityValidator,
|
|
)
|
|
from audita.validators.base import ValidationContext, ValidationDecision, ValidationResult
|
|
|
|
|
|
class RecordingValidator:
|
|
execution_kind = "deterministic"
|
|
|
|
def __init__(self, name, recorder, approve=True):
|
|
self.name = name
|
|
self._recorder = recorder
|
|
self._approve = approve
|
|
|
|
def validate(self, context: ValidationContext) -> ValidationResult:
|
|
self._recorder.append((self.name, [proposal.corrected_text for proposal in context.proposals]))
|
|
return ValidationResult(
|
|
validator_name=self.name,
|
|
execution_kind=self.execution_kind,
|
|
decisions=[
|
|
ValidationDecision(
|
|
proposal_index=proposal.proposal_index,
|
|
approved=self._approve,
|
|
reason=None if self._approve else f"{self.name} rejected proposal",
|
|
)
|
|
for proposal in context.proposals
|
|
],
|
|
)
|
|
|
|
|
|
class RecordingLLMValidator(RecordingValidator):
|
|
execution_kind = "llm"
|
|
|
|
|
|
class FakeStructuredLLMClient:
|
|
def __init__(self, responses):
|
|
self._responses = responses
|
|
self._lock = threading.Lock()
|
|
self.calls = []
|
|
|
|
def run_structured(self, *, stage_name, messages, response_model, config):
|
|
with self._lock:
|
|
self.calls.append(
|
|
{
|
|
"stage_name": stage_name,
|
|
"messages": list(messages),
|
|
"response_model": response_model,
|
|
}
|
|
)
|
|
payload = _pop_llm_response(self._responses, stage_name)
|
|
return response_model.model_validate(payload)
|
|
|
|
|
|
def _pop_llm_response(responses, stage_name):
|
|
if isinstance(responses, dict):
|
|
if stage_name not in responses:
|
|
raise AuditaLLMError(f"FakeStructuredLLMClient received unexpected stage_name: {stage_name}")
|
|
payloads = responses[stage_name]
|
|
if isinstance(payloads, list):
|
|
if not payloads:
|
|
raise AuditaLLMError(f"FakeStructuredLLMClient received too many calls for stage_name: {stage_name}")
|
|
return payloads.pop(0)
|
|
payload = payloads
|
|
del responses[stage_name]
|
|
return payload
|
|
if not responses:
|
|
raise AuditaLLMError("FakeStructuredLLMClient received more calls than expected.")
|
|
return responses.pop(0)
|
|
|
|
|
|
class TrackingStructuredLLMClient:
|
|
def __init__(self, responses, barrier=None):
|
|
self._responses = responses
|
|
self._barrier = barrier
|
|
self._lock = threading.Lock()
|
|
self.calls = []
|
|
self.in_flight = 0
|
|
self.max_in_flight = 0
|
|
|
|
def run_structured(self, *, stage_name, messages, response_model, config):
|
|
if self._barrier is not None:
|
|
self._barrier.wait()
|
|
with self._lock:
|
|
self.in_flight += 1
|
|
self.max_in_flight = max(self.max_in_flight, self.in_flight)
|
|
self.calls.append(
|
|
{
|
|
"stage_name": stage_name,
|
|
"messages": list(messages),
|
|
"response_model": response_model,
|
|
}
|
|
)
|
|
payload = _pop_llm_response(self._responses, stage_name)
|
|
try:
|
|
return response_model.model_validate(payload)
|
|
finally:
|
|
with self._lock:
|
|
self.in_flight -= 1
|
|
|
|
|
|
class RecordingModule:
|
|
replacement_policy = "require_unique"
|
|
|
|
def __init__(self, module_key, proposals, validators, recorder):
|
|
self.module_key = module_key
|
|
self._proposals = proposals
|
|
self._validators = validators
|
|
self._recorder = recorder
|
|
|
|
def validators(self):
|
|
return list(self._validators)
|
|
|
|
def propose(self, transcript_section, context: ModuleContext):
|
|
self._recorder.append(("propose", [item.segment.text for item in transcript_section.segments]))
|
|
return list(self._proposals)
|
|
|
|
|
|
class ConcurrentRecordingModule:
|
|
replacement_policy = "require_unique"
|
|
|
|
def __init__(self, recorder, barrier):
|
|
self.module_key = "concurrent"
|
|
self._recorder = recorder
|
|
self._barrier = barrier
|
|
|
|
def validators(self):
|
|
return []
|
|
|
|
def propose(self, transcript_section, context: ModuleContext):
|
|
texts = [item.segment.text for item in transcript_section.segments]
|
|
self._recorder.append(("start", transcript_section.section_index, texts))
|
|
self._barrier.wait()
|
|
segment = transcript_section.segments[0].segment
|
|
return [
|
|
CorrectionProposal(
|
|
proposal_index=0,
|
|
module_instance=context.run_spec.instance_name,
|
|
module_key=context.run_spec.module_key,
|
|
id=segment.id,
|
|
original_text=segment.text.rstrip("."),
|
|
corrected_text=f"{segment.text.rstrip('.')} revised",
|
|
confidence=0.9,
|
|
)
|
|
]
|
|
|
|
|
|
def test_pipeline_runner_applies_modules_sequentially(tmp_path):
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Alpha."}
|
|
]
|
|
"""
|
|
)
|
|
glossary = parse_glossary_yaml(
|
|
"""
|
|
glossary:
|
|
- name: "Alpha"
|
|
category: noun
|
|
summary: "Alpha."
|
|
"""
|
|
)
|
|
seen = []
|
|
first = RecordingModule(
|
|
"first",
|
|
[
|
|
CorrectionProposal(
|
|
proposal_index=0,
|
|
module_instance="first",
|
|
module_key="first",
|
|
id=1,
|
|
original_text="Alpha",
|
|
corrected_text="Beta",
|
|
confidence=0.9,
|
|
)
|
|
],
|
|
[RecordingValidator("first_validator", seen)],
|
|
seen,
|
|
)
|
|
second = RecordingModule(
|
|
"second",
|
|
[
|
|
CorrectionProposal(
|
|
proposal_index=0,
|
|
module_instance="second",
|
|
module_key="second",
|
|
id=1,
|
|
original_text="Beta",
|
|
corrected_text="Gamma",
|
|
confidence=0.9,
|
|
)
|
|
],
|
|
[RecordingValidator("second_validator", seen)],
|
|
seen,
|
|
)
|
|
|
|
runner = PipelineRunner()
|
|
result = runner.run(
|
|
transcript=transcript,
|
|
glossary=glossary,
|
|
module_specs=[
|
|
ModuleRunSpec(instance_name="first", module_key="first", module=first),
|
|
ModuleRunSpec(instance_name="second", module_key="second", module=second),
|
|
],
|
|
config=AuditaConfig.from_sources(env={}),
|
|
run_dir=tmp_path / "run",
|
|
)
|
|
|
|
assert seen[0] == ("propose", ["Alpha."])
|
|
assert seen[1] == ("first_validator", ["Beta"])
|
|
assert seen[2] == ("propose", ["Beta."])
|
|
assert seen[3] == ("second_validator", ["Gamma"])
|
|
assert result.transcript[0].text == "Gamma."
|
|
assert len(result.applied_changes) == 2
|
|
|
|
|
|
def test_pipeline_runner_validator_order_respects_survivors(tmp_path):
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Hello."}
|
|
]
|
|
"""
|
|
)
|
|
glossary = parse_glossary_yaml(
|
|
"""
|
|
glossary:
|
|
- name: "Hello"
|
|
category: noun
|
|
summary: "Hello."
|
|
"""
|
|
)
|
|
seen = []
|
|
module = RecordingModule(
|
|
"mod",
|
|
[
|
|
CorrectionProposal(
|
|
proposal_index=0,
|
|
module_instance="mod",
|
|
module_key="mod",
|
|
id=1,
|
|
original_text="Hello",
|
|
corrected_text="Goodbye",
|
|
confidence=0.9,
|
|
)
|
|
],
|
|
[
|
|
RecordingValidator("first", seen, approve=False),
|
|
RecordingLLMValidator("second", seen, approve=True),
|
|
],
|
|
seen,
|
|
)
|
|
|
|
runner = PipelineRunner()
|
|
result = runner.run(
|
|
transcript=transcript,
|
|
glossary=glossary,
|
|
module_specs=[ModuleRunSpec(instance_name="mod", module_key="mod", module=module)],
|
|
config=AuditaConfig.from_sources(env={}),
|
|
run_dir=tmp_path / "run",
|
|
)
|
|
|
|
assert ("first", ["Goodbye"]) in seen
|
|
assert all(entry[0] != "second" for entry in seen)
|
|
assert result.module_reports[0].validators[0].rejected_count == 1
|
|
assert result.module_reports[0].validators[1].candidate_count == 0
|
|
assert result.skipped_corrections[0].source == "validator:first"
|
|
|
|
|
|
def test_pipeline_runner_supports_deterministic_and_llm_validators_in_one_chain(tmp_path):
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Alpha."}
|
|
]
|
|
"""
|
|
)
|
|
glossary = parse_glossary_yaml(
|
|
"""
|
|
glossary:
|
|
- name: "Alpha"
|
|
category: noun
|
|
summary: "Alpha."
|
|
"""
|
|
)
|
|
seen = []
|
|
module = RecordingModule(
|
|
"mixed",
|
|
[
|
|
CorrectionProposal(
|
|
proposal_index=0,
|
|
module_instance="mixed",
|
|
module_key="mixed",
|
|
id=1,
|
|
original_text="Alpha",
|
|
corrected_text="Beta",
|
|
confidence=0.9,
|
|
)
|
|
],
|
|
[
|
|
RecordingValidator("deterministic_guard", seen),
|
|
RecordingLLMValidator("llm_review", seen),
|
|
],
|
|
seen,
|
|
)
|
|
|
|
runner = PipelineRunner()
|
|
result = runner.run(
|
|
transcript=transcript,
|
|
glossary=glossary,
|
|
module_specs=[ModuleRunSpec(instance_name="mixed", module_key="mixed", module=module)],
|
|
config=AuditaConfig.from_sources(env={}),
|
|
run_dir=tmp_path / "run",
|
|
)
|
|
|
|
assert result.transcript[0].text == "Beta."
|
|
assert [report.execution_kind for report in result.module_reports[0].validators] == [
|
|
"deterministic",
|
|
"llm",
|
|
]
|
|
|
|
|
|
def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "There were gestures at the dam."}
|
|
]
|
|
"""
|
|
)
|
|
glossary = parse_glossary_yaml(
|
|
"""
|
|
glossary:
|
|
- name: "Jesters"
|
|
category: faction
|
|
summary: "The Jesters are a faction."
|
|
"""
|
|
)
|
|
module = RecordingModule(
|
|
"mixed_real",
|
|
[
|
|
CorrectionProposal(
|
|
proposal_index=0,
|
|
module_instance="mixed_real",
|
|
module_key="glossary",
|
|
id=1,
|
|
original_text="gestures",
|
|
corrected_text="Jesters",
|
|
confidence=0.9,
|
|
)
|
|
],
|
|
[
|
|
ProposalConfidenceValidator("proposal_confidence_guard", "glossary_confidence_threshold"),
|
|
ProtectedGlossaryTermsValidator("protected_glossary_guard"),
|
|
SpokenFormPlausibilityValidator("spoken_form_plausibility_review"),
|
|
MeaningReversalValidator("meaning_reversal_review"),
|
|
],
|
|
[],
|
|
)
|
|
llm_client = FakeStructuredLLMClient(
|
|
{
|
|
"mixed_real:spoken_form_plausibility_review": {
|
|
"validations": [
|
|
{
|
|
"correction_index": 0,
|
|
"approved": True,
|
|
"confidence": 0.95,
|
|
"reason": "Likely phonetic mistranscription in context.",
|
|
}
|
|
]
|
|
},
|
|
"mixed_real:meaning_reversal_review": {
|
|
"validations": [
|
|
{
|
|
"correction_index": 0,
|
|
"approved": True,
|
|
"confidence": 0.98,
|
|
"reason": "Does not reverse the segment meaning.",
|
|
}
|
|
]
|
|
},
|
|
}
|
|
)
|
|
|
|
runner = PipelineRunner()
|
|
result = runner.run(
|
|
transcript=transcript,
|
|
glossary=glossary,
|
|
module_specs=[ModuleRunSpec(instance_name="mixed_real", module_key="glossary", module=module)],
|
|
config=AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "test-key"}),
|
|
run_dir=tmp_path / "run",
|
|
llm_client=llm_client,
|
|
)
|
|
|
|
assert result.transcript[0].text == "There were Jesters at the dam."
|
|
assert [report.execution_kind for report in result.module_reports[0].validators] == [
|
|
"deterministic",
|
|
"deterministic",
|
|
"llm",
|
|
"llm",
|
|
]
|
|
assert {call["stage_name"] for call in llm_client.calls} == {
|
|
"mixed_real:spoken_form_plausibility_review",
|
|
"mixed_real:meaning_reversal_review",
|
|
}
|
|
|
|
|
|
def test_pipeline_runner_uses_real_protected_glossary_validator(tmp_path):
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Hrank moves."}
|
|
]
|
|
"""
|
|
)
|
|
glossary = parse_glossary_yaml(
|
|
"""
|
|
glossary:
|
|
- name: "Hrank"
|
|
category: pc
|
|
summary: "Hrank is a player character."
|
|
"""
|
|
)
|
|
module = RecordingModule(
|
|
"protected",
|
|
[
|
|
CorrectionProposal(
|
|
proposal_index=0,
|
|
module_instance="protected",
|
|
module_key="protected",
|
|
id=1,
|
|
original_text="Hrank",
|
|
corrected_text="Frank",
|
|
confidence=0.9,
|
|
)
|
|
],
|
|
[ProtectedGlossaryTermsValidator("protected_glossary_guard")],
|
|
[],
|
|
)
|
|
|
|
runner = PipelineRunner()
|
|
result = runner.run(
|
|
transcript=transcript,
|
|
glossary=glossary,
|
|
module_specs=[ModuleRunSpec(instance_name="protected", module_key="protected", module=module)],
|
|
config=AuditaConfig.from_sources(env={}),
|
|
run_dir=tmp_path / "run",
|
|
)
|
|
|
|
assert result.transcript[0].text == "Hrank moves."
|
|
assert result.skipped_corrections[0].source == "validator:protected_glossary_guard"
|
|
assert result.skipped_corrections[0].reason == "correction changes protected glossary term usage"
|
|
|
|
|
|
def test_pipeline_runner_collects_section_proposals_concurrently_and_preserves_section_order(tmp_path, monkeypatch):
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Alpha."},
|
|
{"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "Beta."}
|
|
]
|
|
"""
|
|
)
|
|
glossary = parse_glossary_yaml(
|
|
"""
|
|
glossary:
|
|
- name: "Alpha"
|
|
category: noun
|
|
summary: "Alpha."
|
|
"""
|
|
)
|
|
sections = [
|
|
TranscriptSection(
|
|
section_index=0,
|
|
start_index=0,
|
|
segments=[IndexedSegment(index=0, segment=transcript[0])],
|
|
token_count=1,
|
|
),
|
|
TranscriptSection(
|
|
section_index=1,
|
|
start_index=1,
|
|
segments=[IndexedSegment(index=1, segment=transcript[1])],
|
|
token_count=1,
|
|
),
|
|
]
|
|
seen = []
|
|
module = ConcurrentRecordingModule(seen, threading.Barrier(2, timeout=1.0))
|
|
monkeypatch.setattr(
|
|
"audita.framework.runner.chunk_transcript",
|
|
lambda working, max_tokens, min_section_tokens=1, target_section_count=None, exact_target_section_count=None: sections,
|
|
)
|
|
|
|
runner = PipelineRunner()
|
|
result = runner.run(
|
|
transcript=transcript,
|
|
glossary=glossary,
|
|
module_specs=[ModuleRunSpec(instance_name="concurrent", module_key="concurrent", module=module)],
|
|
config=AuditaConfig.from_sources(env={}, overrides=ConfigOverrides(llm_concurrency=2)),
|
|
run_dir=tmp_path / "run",
|
|
)
|
|
|
|
assert [change.id for change in result.applied_changes] == [1, 2]
|
|
assert result.applied_changes[0].corrected_text == "Alpha revised"
|
|
assert result.applied_changes[1].corrected_text == "Beta revised"
|
|
assert result.transcript[0].text == "Alpha revised."
|
|
assert result.transcript[1].text == "Beta revised."
|
|
|
|
|
|
def test_pipeline_runner_passes_exact_target_sections_to_chunker(tmp_path, monkeypatch):
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "Alpha."}
|
|
]
|
|
"""
|
|
)
|
|
glossary = parse_glossary_yaml(
|
|
"""
|
|
glossary:
|
|
- name: "Alpha"
|
|
category: noun
|
|
summary: "Alpha."
|
|
"""
|
|
)
|
|
seen_args = {}
|
|
sections = [
|
|
TranscriptSection(
|
|
section_index=0,
|
|
start_index=0,
|
|
segments=[IndexedSegment(index=0, segment=transcript[0])],
|
|
token_count=1,
|
|
)
|
|
]
|
|
module = RecordingModule("noop", [], [], [])
|
|
|
|
def _fake_chunk_transcript(working, max_tokens, min_section_tokens=1, target_section_count=None, exact_target_section_count=None):
|
|
seen_args["target_section_count"] = target_section_count
|
|
seen_args["exact_target_section_count"] = exact_target_section_count
|
|
return sections
|
|
|
|
monkeypatch.setattr("audita.framework.runner.chunk_transcript", _fake_chunk_transcript)
|
|
|
|
runner = PipelineRunner()
|
|
runner.run(
|
|
transcript=transcript,
|
|
glossary=glossary,
|
|
module_specs=[ModuleRunSpec(instance_name="noop", module_key="noop", module=module)],
|
|
config=AuditaConfig.from_sources(env={}, overrides=ConfigOverrides(target_sections=3)),
|
|
run_dir=tmp_path / "run",
|
|
)
|
|
|
|
assert seen_args == {
|
|
"target_section_count": None,
|
|
"exact_target_section_count": 3,
|
|
}
|
|
|
|
|
|
def test_pipeline_runner_reports_first_llm_validator_rejection_in_chain_order(tmp_path):
|
|
transcript = parse_transcript_json(
|
|
"""
|
|
[
|
|
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "There were gestures at the dam."}
|
|
]
|
|
"""
|
|
)
|
|
glossary = parse_glossary_yaml(
|
|
"""
|
|
glossary:
|
|
- name: "Jesters"
|
|
category: faction
|
|
summary: "The Jesters are a faction."
|
|
"""
|
|
)
|
|
module = RecordingModule(
|
|
"ordered_real",
|
|
[
|
|
CorrectionProposal(
|
|
proposal_index=0,
|
|
module_instance="ordered_real",
|
|
module_key="glossary",
|
|
id=1,
|
|
original_text="gestures",
|
|
corrected_text="Jesters",
|
|
confidence=0.9,
|
|
)
|
|
],
|
|
[
|
|
ProposalConfidenceValidator("proposal_confidence_guard", "glossary_confidence_threshold"),
|
|
ProtectedGlossaryTermsValidator("protected_glossary_guard"),
|
|
SpokenFormPlausibilityValidator("spoken_form_plausibility_review"),
|
|
MeaningReversalValidator("meaning_reversal_review"),
|
|
],
|
|
[],
|
|
)
|
|
llm_client = FakeStructuredLLMClient(
|
|
{
|
|
"ordered_real:spoken_form_plausibility_review": {
|
|
"validations": [
|
|
{
|
|
"correction_index": 0,
|
|
"approved": False,
|
|
"confidence": 0.95,
|
|
"reason": "Not plausibly supported by spoken-form context.",
|
|
}
|
|
]
|
|
},
|
|
"ordered_real:meaning_reversal_review": {
|
|
"validations": [
|
|
{
|
|
"correction_index": 0,
|
|
"approved": False,
|
|
"confidence": 0.98,
|
|
"reason": "Changes meaning too much.",
|
|
}
|
|
]
|
|
},
|
|
}
|
|
)
|
|
|
|
runner = PipelineRunner()
|
|
result = runner.run(
|
|
transcript=transcript,
|
|
glossary=glossary,
|
|
module_specs=[ModuleRunSpec(instance_name="ordered_real", module_key="glossary", module=module)],
|
|
config=AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "test-key"}),
|
|
run_dir=tmp_path / "run",
|
|
llm_client=llm_client,
|
|
)
|
|
|
|
assert result.skipped_corrections[0].source == "validator:spoken_form_plausibility_review"
|
|
assert result.skipped_corrections[0].reason == "Not plausibly supported by spoken-form context."
|