Enhancements to LLM concurrency to improve overall throughput

tests/test_framework_runner.py

@@ -45,23 +45,70 @@ class RecordingLLMValidator(RecordingValidator):
 
 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.")
-        payload = 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)
+
+
+class TrackingStructuredLLMClient:
+    def __init__(self, responses, barrier=None):
+        self._responses = responses
+        self._barrier = barrier
+        self._lock = threading.Lock()
+        self.calls = []
+        self.in_flight = 0
+        self.max_in_flight = 0
+
+    def run_structured(self, *, stage_name, messages, response_model, config):
+        if self._barrier is not None:
+            self._barrier.wait()
+        with self._lock:
+            self.in_flight += 1
+            self.max_in_flight = max(self.max_in_flight, self.in_flight)
+            self.calls.append(
+                {
+                    "stage_name": stage_name,
+                    "messages": list(messages),
+                    "response_model": response_model,
+                }
+            )
+            payload = _pop_llm_response(self._responses, stage_name)
+        try:
+            return response_model.model_validate(payload)
+        finally:
+            with self._lock:
+                self.in_flight -= 1
+
+
 class RecordingModule:
     replacement_policy = "require_unique"
 
@@ -322,8 +369,8 @@ def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
         [],
     )
     llm_client = FakeStructuredLLMClient(
-        [
-            {
+        {
+            "mixed_real:spoken_form_plausibility_review": {
                 "validations": [
                     {
                         "correction_index": 0,
@@ -333,7 +380,7 @@ def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
                     }
                 ]
             },
-            {
+            "mixed_real:meaning_reversal_review": {
                 "validations": [
                     {
                         "correction_index": 0,
@@ -343,7 +390,7 @@ def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
                     }
                 ]
             },
-        ]
+        }
     )
 
     runner = PipelineRunner()
@@ -363,10 +410,10 @@ def test_pipeline_runner_supports_real_llm_validators_in_one_chain(tmp_path):
         "llm",
         "llm",
     ]
-    assert [call["stage_name"] for call in llm_client.calls] == [
+    assert {call["stage_name"] for call in llm_client.calls} == {
         "mixed_real:spoken_form_plausibility_review",
         "mixed_real:meaning_reversal_review",
-    ]
+    }
 
 
 def test_pipeline_runner_uses_real_protected_glossary_validator(tmp_path):
@@ -449,7 +496,10 @@ def test_pipeline_runner_collects_section_proposals_concurrently_and_preserves_s
     ]
     seen = []
     module = ConcurrentRecordingModule(seen, threading.Barrier(2, timeout=1.0))
-    monkeypatch.setattr("audita.framework.runner.chunk_transcript", lambda working, max_tokens: sections)
+    monkeypatch.setattr(
+        "audita.framework.runner.chunk_transcript",
+        lambda working, max_tokens, min_section_tokens=1, target_section_count=None: sections,
+    )
 
     runner = PipelineRunner()
     result = runner.run(
@@ -465,3 +515,79 @@ def test_pipeline_runner_collects_section_proposals_concurrently_and_preserves_s
     assert result.applied_changes[1].corrected_text == "Beta revised"
     assert result.transcript[0].text == "Alpha revised."
     assert result.transcript[1].text == "Beta revised."
+
+
+def test_pipeline_runner_reports_first_llm_validator_rejection_in_chain_order(tmp_path):
+    transcript = parse_transcript_json(
+        """
+        [
+          {"id": 1, "speaker": "A", "start": 0.0, "end": 1.0, "text": "There were gestures at the dam."}
+        ]
+        """
+    )
+    glossary = parse_glossary_yaml(
+        """
+        glossary:
+          - name: "Jesters"
+            category: faction
+            summary: "The Jesters are a faction."
+        """
+    )
+    module = RecordingModule(
+        "ordered_real",
+        [
+            CorrectionProposal(
+                proposal_index=0,
+                module_instance="ordered_real",
+                module_key="glossary",
+                id=1,
+                original_text="gestures",
+                corrected_text="Jesters",
+                confidence=0.9,
+            )
+        ],
+        [
+            ProposalConfidenceValidator("proposal_confidence_guard", "glossary_confidence_threshold"),
+            ProtectedGlossaryTermsValidator("protected_glossary_guard"),
+            SpokenFormPlausibilityValidator("spoken_form_plausibility_review"),
+            MeaningReversalValidator("meaning_reversal_review"),
+        ],
+        [],
+    )
+    llm_client = FakeStructuredLLMClient(
+        {
+            "ordered_real:spoken_form_plausibility_review": {
+                "validations": [
+                    {
+                        "correction_index": 0,
+                        "approved": False,
+                        "confidence": 0.95,
+                        "reason": "Not plausibly supported by spoken-form context.",
+                    }
+                ]
+            },
+            "ordered_real:meaning_reversal_review": {
+                "validations": [
+                    {
+                        "correction_index": 0,
+                        "approved": False,
+                        "confidence": 0.98,
+                        "reason": "Changes meaning too much.",
+                    }
+                ]
+            },
+        }
+    )
+
+    runner = PipelineRunner()
+    result = runner.run(
+        transcript=transcript,
+        glossary=glossary,
+        module_specs=[ModuleRunSpec(instance_name="ordered_real", module_key="glossary", module=module)],
+        config=AuditaConfig.from_sources(env={"OPENROUTER_API_KEY": "test-key"}),
+        run_dir=tmp_path / "run",
+        llm_client=llm_client,
+    )
+
+    assert result.skipped_corrections[0].source == "validator:spoken_form_plausibility_review"
+    assert result.skipped_corrections[0].reason == "Not plausibly supported by spoken-form context."