Implemented optional concurrency for the LLM backend

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