Enhancements to LLM concurrency to improve overall throughput

tests/test_module_proposals.py

@@ -1,3 +1,4 @@
+import threading
 from pathlib import Path
 
 from audita.core.chunking import chunk_transcript
@@ -19,20 +20,38 @@ from audita.pipeline import process_transcript_result
 
 class FakeStructuredLLMClient:
     def __init__(self, responses):
-        self._responses = list(responses)
+        self._responses = responses
+        self._lock = threading.Lock()
         self.calls = []
 
     def run_structured(self, *, stage_name, messages, response_model, config):
-        self.calls.append(
-            {
-                "stage_name": stage_name,
-                "messages": list(messages),
-                "response_model": response_model,
-            }
-        )
-        if not self._responses:
-            raise AuditaLLMError("FakeStructuredLLMClient received more calls than expected.")
-        return response_model.model_validate(self._responses.pop(0))
+        with self._lock:
+            self.calls.append(
+                {
+                    "stage_name": stage_name,
+                    "messages": list(messages),
+                    "response_model": response_model,
+                }
+            )
+            payload = _pop_llm_response(self._responses, stage_name)
+        return response_model.model_validate(payload)
+
+
+def _pop_llm_response(responses, stage_name):
+    if isinstance(responses, dict):
+        if stage_name not in responses:
+            raise AuditaLLMError(f"FakeStructuredLLMClient received unexpected stage_name: {stage_name}")
+        payloads = responses[stage_name]
+        if isinstance(payloads, list):
+            if not payloads:
+                raise AuditaLLMError(f"FakeStructuredLLMClient received too many calls for stage_name: {stage_name}")
+            return payloads.pop(0)
+        payload = payloads
+        del responses[stage_name]
+        return payload
+    if not responses:
+        raise AuditaLLMError("FakeStructuredLLMClient received more calls than expected.")
+    return responses.pop(0)
 
 
 def _glossary():
@@ -810,8 +829,8 @@ def test_process_transcript_result_grammar_module_still_rejects_homophone_style_
         work_dir_retention="always",
     )
     client = FakeStructuredLLMClient(
-        [
-            {
+        {
+            "grammar:proposal": {
                 "corrections": [
                     {
                         "id": 1,
@@ -821,7 +840,7 @@ def test_process_transcript_result_grammar_module_still_rejects_homophone_style_
                     }
                 ]
             },
-            {
+            "grammar:grammar_only_guard": {
                 "validations": [
                     {
                         "correction_index": 0,
@@ -830,8 +849,18 @@ def test_process_transcript_result_grammar_module_still_rejects_homophone_style_
                         "reason": "Free-standing homophone rewrite rather than conservative grammar cleanup.",
                     }
                 ]
-            }
-        ]
+            },
+            "grammar:meaning_reversal_review": {
+                "validations": [
+                    {
+                        "correction_index": 0,
+                        "approved": True,
+                        "confidence": 0.99,
+                        "reason": "Does not reverse the segment meaning.",
+                    }
+                ]
+            },
+        }
     )
 
     result = process_transcript_result(
@@ -846,6 +875,7 @@ def test_process_transcript_result_grammar_module_still_rejects_homophone_style_
     assert [call["stage_name"] for call in client.calls] == [
         "grammar:proposal",
         "grammar:grammar_only_guard",
+        "grammar:meaning_reversal_review",
     ]
     assert result.report.skipped_corrections[0].source == "validator:grammar_only_guard"
     assert "grammar cleanup" in result.report.skipped_corrections[0].reason