Implemented optional concurrency for the LLM backend

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,

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,

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/"

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:
+        section_proposals = _collect_module_proposals(sections=sections, context=context)
-            proposed = list(module.propose(section, 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],

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:
+            if context.config.llm_concurrency == 1 or len(batches) <= 1:
-                payload = [item.to_prompt_payload() for item in batch.items]
+                for batch in batches:
-                messages = self.prompt_builder(payload)
+                    decisions.extend(self._run_batch(context, llm_client, batch))
-                prompt_path = context.run_dir / f"{self.name}-prompt-{batch.batch_index:04d}.json"
+            else:
-                response_path = context.run_dir / f"{self.name}-response-{batch.batch_index:04d}.json"
+                with ThreadPoolExecutor(max_workers=context.config.llm_concurrency) as executor:
-                _write_json(prompt_path, {"messages": messages})
+                    for batch_decisions in executor.map(
-                response = llm_client.run_structured(
+                        lambda batch: self._run_batch(context, llm_client, batch),
-                    stage_name=f"{context.run_spec.instance_name}:{self.name}",
+                        batches,
-                    messages=messages,
+                    ):
-                    response_model=_LLMValidationSetModel,
+                        decisions.extend(batch_decisions)
-                    config=context.config,
-                )
-                _write_json(response_path, response.model_dump(mode="json"))
-                decisions.extend(self._validate_batch_response(response, batch.items))
 
         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):

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"}]

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."

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

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

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