Faster Persp generation

tests/test_statement_personal_flow.py

@@ -2,6 +2,7 @@ import io
 import json
 import os
 import tempfile
+import threading
 import unittest
 from pathlib import Path
 from tkinter import ttk
@@ -95,6 +96,25 @@ class PersonalStatementTests(unittest.TestCase):
 
 
 class SelectiveGeminiTests(unittest.TestCase):
+    def test_rubric_batches_preserve_groups_and_six_question_limit(self):
+        for sizes, expected in (
+            ([2, 2, 2], [6]),
+            ([4, 3, 2, 1], [4, 6]),
+            ([2, 7, 1, 5], [2, 7, 6]),
+            ([], []),
+        ):
+            with self.subTest(sizes=sizes):
+                groups = [[f'{index}:{question}' for question in range(size)]
+                          for index, size in enumerate(sizes)]
+                original = [group[:] for group in groups]
+                batches = gemini.batch_rubric_groups(groups)
+                self.assertEqual([len(batch) for batch in batches], expected)
+                self.assertEqual([label for batch in batches for label in batch],
+                                 [label for group in groups for label in group])
+                for group in groups:
+                    self.assertTrue(any(all(label in batch for label in group) for batch in batches))
+                self.assertEqual(groups, original)
+
     def setUp(self):
         self.temp = tempfile.TemporaryDirectory()
         self.addCleanup(self.temp.cleanup)
@@ -152,15 +172,14 @@ class SelectiveGeminiTests(unittest.TestCase):
 
     def test_rubrics_replace_only_persp_using_normal_prompt(self):
         before = self.snapshot()
-        self.client.models.generate_content.side_effect = [self.rubric_response(self.labels[:2]),
-                                                          self.rubric_response(self.labels[2:])]
+        self.client.models.generate_content.return_value = self.rubric_response(self.labels)
         self.assertEqual(gemini.refine_existing(self.workspace, "persp", api_client=self.client), ExitCode.SUCCESS)
         after = self.snapshot()
         self.assertEqual({key: value for key, value in before.items() if not key.startswith("Persp/")},
                          {key: value for key, value in after.items()
                           if not key.startswith(("Persp/", "Cache/"))})
         responses = list((self.root / 'Cache' / 'rubric_responses').glob('*.json'))
-        self.assertEqual(len(responses), 2)
+        self.assertEqual(len(responses), 1)
         self.assertEqual((self.root / "Persp" / "Ex 1").read_bytes(), before['Persp/Ex 1'])
         for label in self.labels:
             self.assertIn("Barème Gemini", (self.root / "Persp" / label).read_text())
@@ -173,12 +192,34 @@ class SelectiveGeminiTests(unittest.TestCase):
         self.assertIn("caractères mathématiques Unicode", gemini.PROMPT_4)
         self.assertIn(r"$\lfloor \sqrt{k} \rfloor$", gemini.PROMPT_4)
 
+    def separate_requests(self, last):
+        # A full six-question group cannot be combined with the final group.
+        extra = [f'Context {index}' for index in range(4)]
+        self.workspace.labels_file.write_text('\n'.join(self.labels + extra) + '\n')
+        self.workspace.label_groups_file.write_text(
+            ', '.join(self.labels[:2] + extra) + '\n' + self.labels[2] + '\n')
+        for label in extra:
+            for directory in ('Text2', 'Sol2'):
+                (self.root / directory / f'{label}.tex').write_text(label)
+            (self.root / 'Persp' / label).write_text(label + '\nReviewed rubric')
+
+        def response_for_request(**kwargs):
+            text = kwargs['contents'][0].parts[1].text
+            if 'Question [Ex 2]' in text:
+                if isinstance(last, Exception):
+                    raise last
+                return last
+            return self.rubric_response(self.labels[:2])
+
+        self.client.models.generate_content.side_effect = response_for_request
+
     def test_incomplete_or_failed_rubrics_keep_completed_groups(self):
         before = self.snapshot()
         for last in (self.response({"rubrics": []}), RuntimeError("API unavailable")):
             for label in self.labels[:2]:
                 (self.root / 'Persp' / label).write_bytes(before[f'Persp/{label}'])
-            self.client.models.generate_content.side_effect = [self.rubric_response(self.labels[:2]), last]
+            self.separate_requests(last)
+            before = self.snapshot()
             with self.assertRaises((CliError, RuntimeError)):
                 gemini.refine_existing(self.workspace, "persp", api_client=self.client)
             after = self.snapshot()
@@ -186,7 +227,8 @@ class SelectiveGeminiTests(unittest.TestCase):
                 self.assertIn(b'Bar\xc3\xa8me Gemini', after[f'Persp/{label}'])
             self.assertEqual(after['Persp/Ex 2'], before['Persp/Ex 2'])
             self.assertEqual(after['Persp/Ex 1'], before['Persp/Ex 1'])
-            self.assertEqual({key: value for key, value in before.items() if not key.startswith('Persp/')},
+            self.assertEqual({key: value for key, value in before.items()
+                              if not key.startswith(('Persp/', 'Cache/'))},
                              {key: value for key, value in after.items()
                               if not key.startswith(('Persp/', 'Cache/'))})
 
@@ -198,7 +240,7 @@ class SelectiveGeminiTests(unittest.TestCase):
         )
         self.client.models.generate_content.side_effect = [
             self.response({'rubrics': [{'label': label, 'rubric_content': expected} for label in group]})
-            for group in (self.labels[:2], self.labels[2:])
+            for group in (self.labels,)
         ]
         gemini.refine_existing(self.workspace, 'persp', api_client=self.client)
         for label in self.labels:
@@ -220,8 +262,8 @@ class SelectiveGeminiTests(unittest.TestCase):
     def test_control_character_keeps_progress_and_saves_diagnostics(self):
         before = self.snapshot()
         broken = self.response({'rubrics': [{'label': 'Ex 2', 'rubric_content': 'invalid \x01 control'}]})
-        self.client.models.generate_content.side_effect = [
-            self.rubric_response(self.labels[:2]), broken]
+        self.separate_requests(broken)
+        before = self.snapshot()
         output = io.StringIO()
         with patch('sys.stdout', output), self.assertRaisesRegex(ValueError, 'control characters'):
             gemini.refine_existing(self.workspace, 'persp', api_client=self.client)
@@ -258,19 +300,17 @@ class SelectiveGeminiTests(unittest.TestCase):
         before = kept.read_bytes()
         (self.root / 'Persp' / self.labels[1]).write_text('')
         (self.root / 'Persp' / self.labels[2]).unlink()
-        self.client.models.generate_content.side_effect = [
-            self.rubric_response([self.labels[1]]), self.rubric_response([self.labels[2]])]
+        self.client.models.generate_content.return_value = self.rubric_response(self.labels[1:])
         gemini.refine_existing(self.workspace, 'persp', api_client=self.client)
         self.assertEqual(kept.read_bytes(), before)
         call = self.client.models.generate_content.call_args_list[0]
         text = call.kwargs['contents'][0].parts[1].text
-        self.assertIn('UNIQUEMENT pour les labels suivants : ' + self.labels[1] + '\n', text)
+        self.assertIn('UNIQUEMENT pour les labels suivants : ' + ', '.join(self.labels[1:]) + '\n', text)
         self.assertIn('Question [' + self.labels[0] + ']', text)
-        self.assertEqual(self.client.models.generate_content.call_count, 2)
+        self.assertEqual(self.client.models.generate_content.call_count, 1)
 
     def test_retry_requests_only_failed_group(self):
-        self.client.models.generate_content.side_effect = [
-            self.rubric_response(self.labels[:2]), RuntimeError('API unavailable')]
+        self.separate_requests(RuntimeError('API unavailable'))
         with self.assertRaises(RuntimeError):
             gemini.refine_existing(self.workspace, 'persp', api_client=self.client)
         before = {label: (self.root / 'Persp' / label).read_bytes() for label in self.labels[:2]}
@@ -290,13 +330,56 @@ class SelectiveGeminiTests(unittest.TestCase):
             client.assert_not_called()
         self.assertEqual(self.snapshot(), before)
 
+    def test_four_concurrent_requests_save_successes_and_retry_only_failure(self):
+        labels = [f'Q{index}' for index in range(20)]
+        groups = [labels[index:index + 4] for index in range(0, 20, 4)]
+        self.workspace.labels_file.write_text('\n'.join(labels) + '\n')
+        group_text = ''.join(', '.join(group) + '\n' for group in groups)
+        self.workspace.label_groups_file.write_text(group_text)
+        for label in labels:
+            for directory in ('Text2', 'Sol2'):
+                (self.root / directory / f'{label}.tex').write_text(label)
+        barrier = threading.Barrier(4, timeout=10)
+        lock = threading.Lock()
+        active = 0
+        peak = 0
+
+        def generate(client, context, *, workspace, expected_labels):
+            nonlocal active, peak
+            self.assertIn(expected_labels, groups)
+            if expected_labels == groups[-1]:
+                raise RuntimeError('API unavailable')
+            with lock:
+                active += 1
+                peak = max(peak, active)
+            try:
+                barrier.wait()
+                return {label: 'Generated rubric' for label in expected_labels}
+            finally:
+                with lock:
+                    active -= 1
+
+        with patch.object(gemini, 'generate_rubrics', side_effect=generate) as requests:
+            with self.assertRaisesRegex(RuntimeError, 'API unavailable'):
+                gemini.refine_existing(self.workspace, 'persp', api_client=self.client)
+            self.assertEqual(requests.call_count, 5)
+        self.assertEqual(peak, 4)
+        for label in labels[:16]:
+            self.assertEqual((self.root / 'Persp' / label).read_text(), label + '\nGenerated rubric')
+        self.assertFalse((self.root / 'Persp' / labels[-1]).exists())
+        self.client.models.generate_content.return_value = self.rubric_response(groups[-1])
+        gemini.refine_existing(self.workspace, 'persp', api_client=self.client)
+        self.assertEqual(self.client.models.generate_content.call_count, 1)
+        self.assertEqual(self.workspace.label_groups_file.read_text(), group_text)
+        for label in labels:
+            self.assertTrue((self.root / 'Persp' / label).is_file())
+
     def test_delete_persp_forces_regeneration_and_removes_old_files(self):
         for label in self.labels:
             (self.root / 'Persp' / label).write_text(label + '\nReviewed rubric')
-        self.client.models.generate_content.side_effect = [
-            self.rubric_response(self.labels[:2]), self.rubric_response(self.labels[2:])]
+        self.client.models.generate_content.return_value = self.rubric_response(self.labels)
         gemini.refine_existing(self.workspace, 'persp', api_client=self.client, delete_persp=True)
-        self.assertEqual(self.client.models.generate_content.call_count, 2)
+        self.assertEqual(self.client.models.generate_content.call_count, 1)
         self.assertEqual({path.name for path in (self.root / 'Persp').iterdir()}, set(self.labels))
 
     def test_cli_delete_persp_dispatch(self):