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

2
src/audita/cli.py
View File

@@ -28,6 +28,7 @@ def _build_parser() -> argparse.ArgumentParser:
process.add_argument("--output", type=Path, help="write corrected transcript JSON to this path")
process.add_argument("--report-json", type=Path, help="write structured run report JSON to this path")
process.add_argument("--llm-api-key", help="LLM API key for the configured OpenAI-compatible endpoint")
process.add_argument("--llm-concurrency", type=int, help="maximum concurrent LLM calls within a module stage")
process.add_argument(
"--modules",
help="comma-separated module keys to run, for example: grammar or glossary,homophones,glossary",
@@ -90,6 +91,7 @@ def _process(args: argparse.Namespace) -> int:
config = AuditaConfig.from_sources(
overrides=ConfigOverrides(
llm_api_key=args.llm_api_key,
llm_concurrency=args.llm_concurrency,
module_keys=args.modules,
model=args.model,
base_url=args.base_url,

12
src/audita/core/config.py
View File

@@ -11,6 +11,7 @@ from .errors import AuditaConfigError
DEFAULT_MODEL = "openrouter/google/gemma-4-31b-it"
DEFAULT_BASE_URL = "https://openrouter.ai/api/v1"
DEFAULT_LLM_CONCURRENCY = 1
DEFAULT_MAX_RETRIES = 3
DEFAULT_MAX_SECTION_TOKENS = 6144
DEFAULT_GLOSSARY_CONFIDENCE_THRESHOLD = 0.80
@@ -28,6 +29,7 @@ DEFAULT_NORMALIZE_MAX_SEGMENT_TOKENS = 2048
@dataclass(frozen=True)
class ConfigOverrides:
llm_api_key: Optional[str] = None
llm_concurrency: Optional[int] = None
module_keys: Optional[Union[str, Sequence[str]]] = None
model: Optional[str] = None
base_url: Optional[str] = None
@@ -48,6 +50,7 @@ class ConfigOverrides:
@dataclass(frozen=True)
class AuditaConfig:
api_key: Optional[str] = None
llm_concurrency: int = DEFAULT_LLM_CONCURRENCY
module_keys: Tuple[str, ...] = DEFAULT_MODULE_KEYS
model: str = DEFAULT_MODEL
base_url: str = DEFAULT_BASE_URL
@@ -78,6 +81,12 @@ class AuditaConfig:
source.get("AUDITA_LLM_API_KEY"),
source.get("OPENROUTER_API_KEY"),
),
llm_concurrency=_select_int(
selected.llm_concurrency,
source.get("AUDITA_LLM_CONCURRENCY"),
DEFAULT_LLM_CONCURRENCY,
"AUDITA_LLM_CONCURRENCY",
),
module_keys=_select_module_keys(
selected.module_keys,
source.get("AUDITA_MODULES"),
@@ -160,6 +169,8 @@ class AuditaConfig:
def validate(self) -> None:
normalize_module_keys(self.module_keys)
if self.llm_concurrency <= 0:
raise AuditaConfigError("AUDITA_LLM_CONCURRENCY must be greater than zero.")
if not self.model.strip():
raise AuditaConfigError("AUDITA_MODEL must not be empty.")
if not self.base_url.strip():
@@ -200,6 +211,7 @@ class AuditaConfig:
def to_report_dict(self) -> dict:
return {
"api_key_configured": bool(self.api_key),
"llm_concurrency": self.llm_concurrency,
"module_keys": list(self.module_keys),
"model": self.model,
"base_url": self.base_url,

31
src/audita/framework/llm.py
View File

@@ -1,5 +1,6 @@
from __future__ import annotations
import threading
from typing import Any, Optional, Sequence, Tuple
from audita.core.config import AuditaConfig, DEFAULT_BASE_URL
@@ -10,6 +11,7 @@ class OpenAICompatibleStructuredLLMClient:
def __init__(self) -> None:
self._client = None
self._client_identity: Optional[Tuple[str, str]] = None
self._client_lock = threading.Lock()
def run_structured(
self,
@@ -44,22 +46,23 @@ class OpenAICompatibleStructuredLLMClient:
def _get_client(self, config: AuditaConfig) -> Any:
identity = (config.api_key or "", config.base_url)
if self._client is not None and self._client_identity == identity:
with self._client_lock:
if self._client is not None and self._client_identity == identity:
return self._client
try:
import instructor
from openai import OpenAI
except ImportError as exc:
raise AuditaLLMError(
"The LLM dependencies are not installed. Run `uv sync` before using Audita LLM stages."
) from exc
openai_client = OpenAI(api_key=config.api_key, base_url=config.base_url)
self._client = instructor.patch(openai_client, mode=instructor.Mode.TOOLS)
self._client_identity = identity
return self._client
try:
import instructor
from openai import OpenAI
except ImportError as exc:
raise AuditaLLMError(
"The LLM dependencies are not installed. Run `uv sync` before using Audita LLM stages."
) from exc
openai_client = OpenAI(api_key=config.api_key, base_url=config.base_url)
self._client = instructor.patch(openai_client, mode=instructor.Mode.TOOLS)
self._client_identity = identity
return self._client
def _normalize_openrouter_model(model: str) -> str:
prefix = "openrouter/"

18
src/audita/framework/runner.py
View File

@@ -1,3 +1,4 @@
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
@@ -148,8 +149,8 @@ def _run_module(
applied_changes: List[AppliedChange] = []
updated_transcript = list(working)
try:
for section in sections:
proposed = list(module.propose(section, context))
section_proposals = _collect_module_proposals(sections=sections, context=context)
for proposed in section_proposals:
raw_proposals.extend(proposed)
proposals = [
@@ -258,6 +259,19 @@ def _run_module(
) from exc
def _collect_module_proposals(
*,
sections: Sequence[TranscriptSection],
context: ModuleContext,
) -> List[List[CorrectionProposal]]:
module = context.run_spec.module
if context.config.llm_concurrency == 1 or len(sections) <= 1:
return [list(module.propose(section, context)) for section in sections]
with ThreadPoolExecutor(max_workers=context.config.llm_concurrency) as executor:
return list(executor.map(lambda section: list(module.propose(section, context)), sections))
def _index_validation_decisions(
result: ValidationResult,
proposals: Sequence[CorrectionProposal],

45
src/audita/validators/llm.py
View File

@@ -1,5 +1,6 @@
from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor
import json
from dataclasses import dataclass
from pathlib import Path
@@ -73,20 +74,16 @@ class _BaseLLMValidator:
payload_fn=lambda item: item.to_prompt_payload(),
empty_error_message="Validation input must contain at least one proposal.",
)
for batch in batches:
payload = [item.to_prompt_payload() for item in batch.items]
messages = self.prompt_builder(payload)
prompt_path = context.run_dir / f"{self.name}-prompt-{batch.batch_index:04d}.json"
response_path = context.run_dir / f"{self.name}-response-{batch.batch_index:04d}.json"
_write_json(prompt_path, {"messages": messages})
response = llm_client.run_structured(
stage_name=f"{context.run_spec.instance_name}:{self.name}",
messages=messages,
response_model=_LLMValidationSetModel,
config=context.config,
)
_write_json(response_path, response.model_dump(mode="json"))
decisions.extend(self._validate_batch_response(response, batch.items))
if context.config.llm_concurrency == 1 or len(batches) <= 1:
for batch in batches:
decisions.extend(self._run_batch(context, llm_client, batch))
else:
with ThreadPoolExecutor(max_workers=context.config.llm_concurrency) as executor:
for batch_decisions in executor.map(
lambda batch: self._run_batch(context, llm_client, batch),
batches,
):
decisions.extend(batch_decisions)
return ValidationResult(
validator_name=self.name,
@@ -128,6 +125,26 @@ class _BaseLLMValidator:
for proposal in proposals
]
def _run_batch(
self,
context: ValidationContext,
llm_client,
batch,
) -> List[ValidationDecision]:
payload = [item.to_prompt_payload() for item in batch.items]
messages = self.prompt_builder(payload)
prompt_path = context.run_dir / f"{self.name}-prompt-{batch.batch_index:04d}.json"
response_path = context.run_dir / f"{self.name}-response-{batch.batch_index:04d}.json"
_write_json(prompt_path, {"messages": messages})
response = llm_client.run_structured(
stage_name=f"{context.run_spec.instance_name}:{self.name}",
messages=messages,
response_model=_LLMValidationSetModel,
config=context.config,
)
_write_json(response_path, response.model_dump(mode="json"))
return self._validate_batch_response(response, batch.items)
@dataclass(frozen=True)
class SpokenFormPlausibilityValidator(_BaseLLMValidator):

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