Move python implementation under python/ in preparation for the upcoming Go rewrite
This commit is contained in:
642
python/tests/test_framework_runner.py
Normal file
642
python/tests/test_framework_runner.py
Normal file
@@ -0,0 +1,642 @@
|
||||
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."
|
||||
Reference in New Issue
Block a user