Enhancements to LLM concurrency to improve overall throughput

tests/test_new_pipeline.py

@@ -1,4 +1,5 @@
 import json
+import threading
 
 import pytest
 
@@ -12,25 +13,42 @@ from audita.pipeline import process_transcript, 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.")
-        response = self._responses.pop(0)
+        with self._lock:
+            self.calls.append(
+                {
+                    "stage_name": stage_name,
+                    "messages": list(messages),
+                    "response_model": response_model,
+                }
+            )
+            response = _pop_llm_response(self._responses, stage_name)
         if isinstance(response, Exception):
             raise response
         return response_model.model_validate(response)
 
 
+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():
     return parse_glossary_yaml(
         """