From bbbef37d9a2cfcf48b69bbcaa985e114ddde2a56 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Tue, 28 Apr 2026 22:14:25 -0500 Subject: [PATCH] Implemented optional concurrency for the LLM backend --- src/audita/cli.py | 2 + src/audita/core/config.py | 12 ++++ src/audita/framework/llm.py | 31 +++++---- src/audita/framework/runner.py | 18 +++++- src/audita/validators/llm.py | 45 +++++++++---- tests/test_framework_llm.py | 23 +++++++ tests/test_framework_runner.py | 85 +++++++++++++++++++++++- tests/test_llm_validators.py | 115 ++++++++++++++++++++++++++++++++- tests/test_new_cli.py | 55 ++++++++++++++++ tests/test_new_config.py | 23 +++++++ 10 files changed, 376 insertions(+), 33 deletions(-) diff --git a/src/audita/cli.py b/src/audita/cli.py index d3398e3..8600899 100644 --- a/src/audita/cli.py +++ b/src/audita/cli.py @@ -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, diff --git a/src/audita/core/config.py b/src/audita/core/config.py index 32d6850..1c2882e 100644 --- a/src/audita/core/config.py +++ b/src/audita/core/config.py @@ -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, diff --git a/src/audita/framework/llm.py b/src/audita/framework/llm.py index ae8faf4..034d80c 100644 --- a/src/audita/framework/llm.py +++ b/src/audita/framework/llm.py @@ -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/" diff --git a/src/audita/framework/runner.py b/src/audita/framework/runner.py index 0f4ae0c..69c728c 100644 --- a/src/audita/framework/runner.py +++ b/src/audita/framework/runner.py @@ -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], diff --git a/src/audita/validators/llm.py b/src/audita/validators/llm.py index e9e10c7..3c96390 100644 --- a/src/audita/validators/llm.py +++ b/src/audita/validators/llm.py @@ -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): diff --git a/tests/test_framework_llm.py b/tests/test_framework_llm.py index 8f094ab..c828cb8 100644 --- a/tests/test_framework_llm.py +++ b/tests/test_framework_llm.py @@ -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"}] diff --git a/tests/test_framework_runner.py b/tests/test_framework_runner.py index 49b39f8..151a862 100644 --- a/tests/test_framework_runner.py +++ b/tests/test_framework_runner.py @@ -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." diff --git a/tests/test_llm_validators.py b/tests/test_llm_validators.py index 2fa8056..a4dee2e 100644 --- a/tests/test_llm_validators.py +++ b/tests/test_llm_validators.py @@ -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 diff --git a/tests/test_new_cli.py b/tests/test_new_cli.py index f7284e7..1d69d53 100644 --- a/tests/test_new_cli.py +++ b/tests/test_new_cli.py @@ -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 diff --git a/tests/test_new_config.py b/tests/test_new_config.py index 07dff15..b82ba5f 100644 --- a/tests/test_new_config.py +++ b/tests/test_new_config.py @@ -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})