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

Implemented optional concurrency for the LLM backend

Browse Source
This commit is contained in:
Eric Rakestraw
2026-04-28 22:14:25 -05:00
parent f834ffad97
commit bbbef37d9a
10 changed files with 376 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
tests/test_framework_llm.py
View File

@@ -1,3 +1,4 @@
from concurrent.futures import ThreadPoolExecutor
import sys
import types
@@ -16,6 +17,7 @@ def _config(**overrides):
base = AuditaConfig.from_sources(env={})
data = {
"api_key": "test-key",
"llm_concurrency": base.llm_concurrency,
"module_keys": base.module_keys,
"model": base.model,
"base_url": base.base_url,
@@ -172,3 +174,24 @@ def test_missing_api_key_error_is_provider_neutral():
response_model=DummyResponseModel,
config=_config(api_key=None),
)
def test_client_initialization_is_safe_under_concurrent_calls(monkeypatch):
_, openai_inits = _install_fake_llm_modules(monkeypatch)
client = OpenAICompatibleStructuredLLMClient()
config = _config(api_key="key-1", base_url="http://localhost:8000/v1")
with ThreadPoolExecutor(max_workers=4) as executor:
list(
executor.map(
lambda _: client.run_structured(
stage_name="test-stage",
messages=[{"role": "user", "content": "Hello"}],
response_model=DummyResponseModel,
config=config,
),
range(4),
)
)
assert openai_inits == [{"api_key": "key-1", "base_url": "http://localhost:8000/v1"}]

85
tests/test_framework_runner.py
View File

@@ -1,4 +1,7 @@
from audita.core.config import AuditaConfig
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
@@ -76,6 +79,35 @@ class RecordingModule:
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(
"""
@@ -382,3 +414,54 @@ def test_pipeline_runner_uses_real_protected_glossary_validator(tmp_path):
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: 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."

115
tests/test_llm_validators.py
View File

@@ -1,7 +1,9 @@
import threading
import pytest
from audita.core.chunking import TokenBatch
from audita.core.config import AuditaConfig
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, ModuleRunSpec
@@ -44,6 +46,31 @@ class FakeStructuredLLMClient:
return response_model.model_validate(self._responses.pop(0))
class CoordinatedStructuredLLMClient:
def __init__(self, responses, barrier):
self._responses = list(responses)
self._barrier = barrier
self._lock = threading.Lock()
self.calls = []
def run_structured(self, *, stage_name, messages, response_model, config):
self._barrier.wait()
with self._lock:
self.calls.append(
{
"stage_name": stage_name,
"messages": list(messages),
"response_model": response_model,
}
)
if not self._responses:
raise AuditaLLMError("CoordinatedStructuredLLMClient received more calls than expected.")
payload = self._responses.pop(0)
if callable(payload):
payload = payload(stage_name=stage_name, messages=messages, response_model=response_model, config=config)
return response_model.model_validate(payload)
def _glossary():
return parse_glossary_yaml(
"""
@@ -68,12 +95,16 @@ def _context(
llm_client,
tmp_path,
replacement_policy="require_unique",
llm_concurrency=1,
):
return ValidationContext(
proposals=proposals,
transcript=transcript,
glossary=_glossary(),
config=AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "test-key"}),
config=AuditaConfig.from_sources(
env={"OPENROUTER_API_KEY": "test-key"},
overrides=ConfigOverrides(llm_concurrency=llm_concurrency),
),
run_spec=ModuleRunSpec(
instance_name="homophones",
module_key="homophones",
@@ -929,3 +960,83 @@ def test_llm_validators_use_shared_token_batching_helper(monkeypatch, tmp_path):
assert len(chunk_calls) == 1
assert chunk_calls[0]["max_tokens"] == AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "test-key"}).max_section_tokens
assert all("corrected_segment_text" in payload for payload in chunk_calls[0]["payloads"])
def test_llm_validators_process_batches_concurrently_and_preserve_proposal_order(monkeypatch, tmp_path):
transcript = parse_transcript_json(
"""
[
{"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "There were gestures at the temple."},
{"id": 2, "speaker": "A", "start": 1.0, "end": 2.0, "text": "ChatGPT still can't do that with a dam."}
]
"""
)
proposals = [
CorrectionProposal(
proposal_index=0,
module_instance="homophones",
module_key="homophones",
id=1,
original_text="gestures",
corrected_text="Jesters",
confidence=0.95,
),
CorrectionProposal(
proposal_index=1,
module_instance="homophones",
module_key="homophones",
id=2,
original_text="dam",
corrected_text="damn",
confidence=0.95,
),
]
client = CoordinatedStructuredLLMClient(
[
lambda **kwargs: {
"validations": [
{
"correction_index": 0 if '"correction_index": 0' in kwargs["messages"][1]["content"] else 1,
"approved": True,
"confidence": 0.97,
"reason": "ok",
}
]
},
lambda **kwargs: {
"validations": [
{
"correction_index": 0 if '"correction_index": 0' in kwargs["messages"][1]["content"] else 1,
"approved": True,
"confidence": 0.94,
"reason": "ok",
}
]
},
],
threading.Barrier(2, timeout=1.0),
)
def fake_chunk_payload_items(items, max_tokens, payload_fn, empty_error_message):
return [
TokenBatch(batch_index=0, items=list(items[:1]), token_count=1),
TokenBatch(batch_index=1, items=list(items[1:]), token_count=1),
]
monkeypatch.setattr(llm_module, "chunk_payload_items", fake_chunk_payload_items)
result = SpokenFormPlausibilityValidator("spoken_form_plausibility_review").validate(
_context(
proposals=proposals,
transcript=transcript,
llm_client=client,
tmp_path=tmp_path,
llm_concurrency=2,
)
)
assert [(decision.proposal_index, decision.approved) for decision in result.decisions] == [
(0, True),
(1, True),
]
assert len(client.calls) == 2

55
tests/test_new_cli.py
View File

@@ -21,6 +21,7 @@ def test_process_help_exposes_framework_flags(capsys):
output = capsys.readouterr().out
assert "--report-json" in output
assert "--llm-api-key" in output
assert "--llm-concurrency" in output
assert "--modules" in output
assert "--model" in output
assert "--base-url" in output
@@ -194,3 +195,57 @@ def test_cli_process_passes_llm_api_key_override_to_config(monkeypatch, tmp_path
assert exit_code == 0
assert captured["llm_api_key"] == "cli-key"
def test_cli_process_passes_llm_concurrency_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["llm_concurrency"] = overrides.llm_concurrency
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",
"--llm-concurrency",
"3",
]
)
assert exit_code == 0
assert captured["llm_concurrency"] == 3

23
tests/test_new_config.py
View File

@@ -6,6 +6,7 @@ from audita.core.config import (
DEFAULT_GLOSSARY_CONFIDENCE_THRESHOLD,
DEFAULT_GRAMMAR_CONFIDENCE_THRESHOLD,
DEFAULT_HOMOPHONES_CONFIDENCE_THRESHOLD,
DEFAULT_LLM_CONCURRENCY,
DEFAULT_NORMALIZE_MAX_SEGMENT_GAP,
DEFAULT_SPOKEN_WORD_CONFIDENCE_THRESHOLD,
DEFAULT_WORK_DIR_RETENTION,
@@ -18,6 +19,7 @@ def test_default_config_allows_missing_api_key():
config = AuditaConfig.from_sources(env={})
assert config.api_key is None
assert config.llm_concurrency == DEFAULT_LLM_CONCURRENCY
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
@@ -36,6 +38,21 @@ def test_cli_overrides_take_precedence():
assert config.max_section_tokens == 2000
def test_llm_concurrency_cli_override_takes_precedence():
config = AuditaConfig.from_sources(
env={"AUDITA_LLM_CONCURRENCY": "2"},
overrides=ConfigOverrides(llm_concurrency=4),
)
assert config.llm_concurrency == 4
def test_llm_concurrency_env_is_parsed():
config = AuditaConfig.from_sources(env={"AUDITA_LLM_CONCURRENCY": "3"})
assert config.llm_concurrency == 3
def test_generic_llm_api_key_env_is_read():
config = AuditaConfig.from_sources(env={"AUDITA_LLM_API_KEY": "generic-key"})
@@ -141,3 +158,9 @@ def test_invalid_module_sequences_are_rejected(value):
def test_invalid_thresholds_are_rejected(env_name):
with pytest.raises(AuditaConfigError):
AuditaConfig.from_sources(env={env_name: "1.5"})
@pytest.mark.parametrize("value", ["0", "-1", "many"])
def test_invalid_llm_concurrency_is_rejected(value):
with pytest.raises(AuditaConfigError, match="AUDITA_LLM_CONCURRENCY"):
AuditaConfig.from_sources(env={"AUDITA_LLM_CONCURRENCY": value})