eric/audita
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 2 Wiki Activity

Enhancements to LLM concurrency to improve overall throughput

Browse Source
This commit is contained in:
Eric Rakestraw
2026-04-29 16:04:36 -05:00
parent db6004fadc
commit d8bc84934e
16 changed files with 815 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
tests/test_framework_chunking.py
View File

@@ -27,6 +27,79 @@ def test_chunk_transcript_batches_sections_by_token_limit():
assert [segment.segment.id for segment in sections[1].segments] == [3]
def test_chunk_transcript_targets_llm_concurrency_when_feasible():
transcript = parse_transcript_json(
"""
[
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
{"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
{"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"},
{"id": 4, "speaker": "A", "start": 3.0, "end": 4.0, "text": "four"}
]
"""
)
sections = chunk_transcript(
transcript,
max_section_tokens=8,
min_section_tokens=4,
target_section_count=2,
estimator=FakeEstimator(),
)
assert len(sections) == 2
assert [segment.segment.id for segment in sections[0].segments] == [1, 2]
assert [segment.segment.id for segment in sections[1].segments] == [3, 4]
def test_chunk_transcript_increases_section_count_when_target_sections_exceed_max_tokens():
transcript = parse_transcript_json(
"""
[
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
{"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
{"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"},
{"id": 4, "speaker": "A", "start": 3.0, "end": 4.0, "text": "four"},
{"id": 5, "speaker": "A", "start": 4.0, "end": 5.0, "text": "five"}
]
"""
)
sections = chunk_transcript(
transcript,
max_section_tokens=8,
min_section_tokens=4,
target_section_count=1,
estimator=FakeEstimator(),
)
assert len(sections) > 1
assert all(section.token_count <= 8 for section in sections)
def test_chunk_transcript_reduces_section_count_when_target_sections_fall_below_min_tokens():
transcript = parse_transcript_json(
"""
[
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "one"},
{"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "two"},
{"id": 3, "speaker": "A", "start": 2.0, "end": 3.0, "text": "three"}
]
"""
)
sections = chunk_transcript(
transcript,
max_section_tokens=12,
min_section_tokens=8,
target_section_count=3,
estimator=FakeEstimator(),
)
assert len(sections) == 1
assert sections[0].token_count >= 8
def test_chunk_transcript_prompt_payload_includes_categories_when_present():
transcript = parse_transcript_json(
"""

162
tests/test_framework_runner.py
View File

@@ -45,23 +45,70 @@ class RecordingLLMValidator(RecordingValidator):
class FakeStructuredLLMClient:
def __init__(self, responses):
self._responses = list(responses)
self._responses = responses
self._lock = threading.Lock()
self.calls = []
def run_structured(self, *, stage_name, messages, response_model, config):
self.calls.append(
{
"stage_name": stage_name,
"messages": list(messages),
"response_model": response_model,
}
)
if not self._responses:
raise AuditaLLMError("FakeStructuredLLMClient received more calls than expected.")
payload = self._responses.pop(0)
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"
@@ -322,8 +369,8 @@ def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
[],
)
llm_client = FakeStructuredLLMClient(
[
{
{
"mixed_real:spoken_form_plausibility_review": {
"validations": [
{
"correction_index": 0,
@@ -333,7 +380,7 @@ def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
}
]
},
{
"mixed_real:meaning_reversal_review": {
"validations": [
{
"correction_index": 0,
@@ -343,7 +390,7 @@ def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
}
]
},
]
}
)
runner = PipelineRunner()
@@ -363,10 +410,10 @@ def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
"llm",
"llm",
]
assert [call["stage_name"] for call in llm_client.calls] == [
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):
@@ -449,7 +496,10 @@ def test_pipeline_runner_collects_section_proposals_concurrently_and_preserves_s
]
seen = []
module = ConcurrentRecordingModule(seen, threading.Barrier(2, timeout=1.0))
monkeypatch.setattr("audita.framework.runner.chunk_transcript", lambda working, max_tokens: sections)
monkeypatch.setattr(
"audita.framework.runner.chunk_transcript",
lambda working, max_tokens, min_section_tokens=1, target_section_count=None: sections,
)
runner = PipelineRunner()
result = runner.run(
@@ -465,3 +515,79 @@ def test_pipeline_runner_collects_section_proposals_concurrently_and_preserves_s
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_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."

62
tests/test_module_proposals.py
View File

@@ -1,3 +1,4 @@
import threading
from pathlib import Path
from audita.core.chunking import chunk_transcript
@@ -19,20 +20,38 @@ from audita.pipeline import process_transcript_result
class FakeStructuredLLMClient:
def __init__(self, responses):
self._responses = list(responses)
self._responses = responses
self._lock = threading.Lock()
self.calls = []
def run_structured(self, *, stage_name, messages, response_model, config):
self.calls.append(
{
"stage_name": stage_name,
"messages": list(messages),
"response_model": response_model,
}
)
if not self._responses:
raise AuditaLLMError("FakeStructuredLLMClient received more calls than expected.")
return response_model.model_validate(self._responses.pop(0))
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)
def _glossary():
@@ -810,8 +829,8 @@ def test_process_transcript_result_grammar_module_still_rejects_homophone_style_
work_dir_retention="always",
)
client = FakeStructuredLLMClient(
[
{
{
"grammar:proposal": {
"corrections": [
{
"id": 1,
@@ -821,7 +840,7 @@ def test_process_transcript_result_grammar_module_still_rejects_homophone_style_
}
]
},
{
"grammar:grammar_only_guard": {
"validations": [
{
"correction_index": 0,
@@ -830,8 +849,18 @@ def test_process_transcript_result_grammar_module_still_rejects_homophone_style_
"reason": "Free-standing homophone rewrite rather than conservative grammar cleanup.",
}
]
}
]
},
"grammar:meaning_reversal_review": {
"validations": [
{
"correction_index": 0,
"approved": True,
"confidence": 0.99,
"reason": "Does not reverse the segment meaning.",
}
]
},
}
)
result = process_transcript_result(
@@ -846,6 +875,7 @@ def test_process_transcript_result_grammar_module_still_rejects_homophone_style_
assert [call["stage_name"] for call in client.calls] == [
"grammar:proposal",
"grammar:grammar_only_guard",
"grammar:meaning_reversal_review",
]
assert result.report.skipped_corrections[0].source == "validator:grammar_only_guard"
assert "grammar cleanup" in result.report.skipped_corrections[0].reason

55
tests/test_new_cli.py
View File

@@ -27,6 +27,7 @@ def test_process_help_exposes_framework_flags(capsys):
assert "--base-url" in output
assert "--max-retries" in output
assert "--max-section-tokens" in output
assert "--min-section-tokens" in output
assert "--glossary-confidence-threshold" in output
assert "--grammar-confidence-threshold" in output
assert "--homophones-confidence-threshold" in output
@@ -249,3 +250,57 @@ def test_cli_process_passes_llm_concurrency_override_to_config(monkeypatch, tmp_
assert exit_code == 0
assert captured["llm_concurrency"] == 3
def test_cli_process_passes_min_section_tokens_override_to_config(monkeypatch, tmp_path):
captured = {}
transcript = parse_transcript_json(
"""
[
{"id": 1, "speaker": "Eric", "start": 0.0, "end": 1.0, "text": "Fixed."}
]
"""
)
report = RunReport(
status="success",
config={"model": "m", "base_url": "b"},
normalization={"source_segment_count": 1, "normalized_segment_count": 1, "merge_count": 0},
pipeline=["grammar"],
modules=[],
applied_changes=[],
skipped_corrections=[],
totals={"output_segment_count": 1, "applied_change_count": 0, "skipped_correction_count": 0},
work_dir_retention="auto",
work_dir_retained=False,
work_dir=None,
error=None,
)
result = ProcessResult(
transcript=transcript,
report=report,
run_dir=tmp_path / "run",
work_dir_retained=False,
)
def _fake_from_sources(*, overrides=None):
captured["min_section_tokens"] = overrides.min_section_tokens
return object()
monkeypatch.setattr("audita.cli.AuditaConfig.from_sources", _fake_from_sources)
monkeypatch.setattr("audita.cli.load_transcript", lambda path: [])
monkeypatch.setattr("audita.cli.load_glossary", lambda path: object())
monkeypatch.setattr("audita.cli.process_transcript_result", lambda *args, **kwargs: result)
exit_code = main(
[
"process",
"transcript.json",
"--glossary",
"glossary.yaml",
"--min-section-tokens",
"5000",
]
)
assert exit_code == 0
assert captured["min_section_tokens"] == 5000

37
tests/test_new_config.py
View File

@@ -7,6 +7,8 @@ from audita.core.config import (
DEFAULT_GRAMMAR_CONFIDENCE_THRESHOLD,
DEFAULT_HOMOPHONES_CONFIDENCE_THRESHOLD,
DEFAULT_LLM_CONCURRENCY,
DEFAULT_MAX_SECTION_TOKENS,
DEFAULT_MIN_SECTION_TOKENS,
DEFAULT_NORMALIZE_MAX_SEGMENT_GAP,
DEFAULT_SPOKEN_WORD_CONFIDENCE_THRESHOLD,
DEFAULT_WORK_DIR_RETENTION,
@@ -20,6 +22,8 @@ def test_default_config_allows_missing_api_key():
assert config.api_key is None
assert config.llm_concurrency == DEFAULT_LLM_CONCURRENCY
assert config.max_section_tokens == DEFAULT_MAX_SECTION_TOKENS
assert config.min_section_tokens == DEFAULT_MIN_SECTION_TOKENS
assert config.module_keys == DEFAULT_MODULE_KEYS
assert config.glossary_confidence_threshold == DEFAULT_GLOSSARY_CONFIDENCE_THRESHOLD
assert config.grammar_confidence_threshold == DEFAULT_GRAMMAR_CONFIDENCE_THRESHOLD
@@ -32,12 +36,21 @@ def test_default_config_allows_missing_api_key():
def test_cli_overrides_take_precedence():
config = AuditaConfig.from_sources(
env={"AUDITA_MAX_SECTION_TOKENS": "1000"},
overrides=ConfigOverrides(max_section_tokens=2000),
overrides=ConfigOverrides(max_section_tokens=2000, min_section_tokens=1000),
)
assert config.max_section_tokens == 2000
def test_min_section_tokens_cli_override_takes_precedence():
config = AuditaConfig.from_sources(
env={"AUDITA_MIN_SECTION_TOKENS": "2000"},
overrides=ConfigOverrides(min_section_tokens=6000),
)
assert config.min_section_tokens == 6000
def test_llm_concurrency_cli_override_takes_precedence():
config = AuditaConfig.from_sources(
env={"AUDITA_LLM_CONCURRENCY": "2"},
@@ -53,6 +66,12 @@ def test_llm_concurrency_env_is_parsed():
assert config.llm_concurrency == 3
def test_min_section_tokens_env_is_parsed():
config = AuditaConfig.from_sources(env={"AUDITA_MIN_SECTION_TOKENS": "3000"})
assert config.min_section_tokens == 3000
def test_generic_llm_api_key_env_is_read():
config = AuditaConfig.from_sources(env={"AUDITA_LLM_API_KEY": "generic-key"})
@@ -164,3 +183,19 @@ def test_invalid_thresholds_are_rejected(env_name):
def test_invalid_llm_concurrency_is_rejected(value):
with pytest.raises(AuditaConfigError, match="AUDITA_LLM_CONCURRENCY"):
AuditaConfig.from_sources(env={"AUDITA_LLM_CONCURRENCY": value})
@pytest.mark.parametrize("value", ["0", "-1", "many"])
def test_invalid_min_section_tokens_is_rejected(value):
with pytest.raises(AuditaConfigError, match="AUDITA_MIN_SECTION_TOKENS"):
AuditaConfig.from_sources(env={"AUDITA_MIN_SECTION_TOKENS": value})
def test_min_section_tokens_must_not_exceed_max_section_tokens():
with pytest.raises(AuditaConfigError, match="AUDITA_MIN_SECTION_TOKENS"):
AuditaConfig.from_sources(
env={
"AUDITA_MIN_SECTION_TOKENS": "9000",
"AUDITA_MAX_SECTION_TOKENS": "8000",
}
)

40
tests/test_new_pipeline.py
View File

@@ -1,4 +1,5 @@
import json
import threading
import pytest
@@ -12,25 +13,42 @@ from audita.pipeline import process_transcript, process_transcript_result
class FakeStructuredLLMClient:
def __init__(self, responses):
self._responses = list(responses)
self._responses = responses
self._lock = threading.Lock()
self.calls = []
def run_structured(self, *, stage_name, messages, response_model, config):
self.calls.append(
{
"stage_name": stage_name,
"messages": list(messages),
"response_model": response_model,
}
)
if not self._responses:
raise AuditaLLMError("FakeStructuredLLMClient received more calls than expected.")
response = self._responses.pop(0)
with self._lock:
self.calls.append(
{
"stage_name": stage_name,
"messages": list(messages),
"response_model": response_model,
}
)
response = _pop_llm_response(self._responses, stage_name)
if isinstance(response, Exception):
raise response
return response_model.model_validate(response)
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)
def _glossary():
return parse_glossary_yaml(
"""