обновлен формат генерации датасета

tests/test_assembly_pipeline.py

@@ -0,0 +1,113 @@
+import asyncio
+import json
+from pathlib import Path
+import tempfile
+from types import SimpleNamespace
+import unittest
+from unittest.mock import patch
+
+from assembly_dataset.config import Config, TeacherConfig
+from assembly_dataset.pipeline import generate, export, validate
+from assembly_dataset.teacher import Plan
+from assembly_dataset.world import solve
+from assembly_dataset.language import action_text
+from llm.config import OpenAiConfig
+
+
+class TestAgent:
+    def __init__(self, invalid=False):
+        self.active = 0
+        self.peak = 0
+        self.invalid = invalid
+
+    async def run(self, prompt, **kwargs):
+        self.active += 1
+        self.peak = max(self.peak, self.active)
+        try:
+            await asyncio.sleep(0.001)
+            plan = {'status': 'ready', 'actions': []} if self.invalid else solve(json.loads(prompt)['scene'])[0]
+            if plan['actions']:
+                scene = json.loads(prompt)['scene']
+                plan['actions'] = [{**a, 'text': action_text(a, scene, Config().prompts)} for a in plan['actions']]
+            return SimpleNamespace(output=Plan.model_validate(plan))
+        finally:
+            self.active -= 1
+
+
+class PipelineTests(unittest.TestCase):
+    def teacher_config(self, path):
+        return Config(output_dir=Path(path), target_examples=30,
+                      teacher=TeacherConfig(enabled=True, fraction=1, concurrency=2, retries=1,
+                                            llm=OpenAiConfig(provider='openai', model='test')))
+
+    def test_teacher_concurrency_and_validated_provenance(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            config = self.teacher_config(tmp)
+            agent = TestAgent()
+            asyncio.run(generate(config, agent))
+            report = export(config)
+            self.assertEqual(agent.peak, 2)
+            self.assertEqual(report['sources'], {'openai': 10})
+            self.assertEqual(report['teacher_failures'], 0)
+            self.assertEqual(validate(config)['examples_checked'], 30)
+
+    def test_invalid_teacher_cannot_poison_corpus(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            config = self.teacher_config(tmp)
+            asyncio.run(generate(config, TestAgent(invalid=True)))
+            report = export(config)
+            self.assertEqual(report['teacher_failures'], 10)
+            self.assertEqual(report['sources'], {'symbolic': 10})
+            self.assertEqual(validate(config)['examples_checked'], 30)
+
+    def test_restart_after_interrupted_generation_equals_clean_run(self):
+        import assembly_dataset.pipeline as pipeline
+        original = pipeline.make_episode
+        def interrupted(index, config):
+            if index >= 40:
+                raise KeyboardInterrupt()
+            return original(index, config)
+        with tempfile.TemporaryDirectory() as first, tempfile.TemporaryDirectory() as second:
+            config = Config(output_dir=Path(first), target_examples=150)
+            with patch.object(pipeline, 'make_episode', side_effect=interrupted):
+                with self.assertRaises(KeyboardInterrupt):
+                    asyncio.run(generate(config))
+            asyncio.run(generate(config))
+            resumed = export(config)
+            clean = Config(output_dir=Path(second), target_examples=150)
+            asyncio.run(generate(clean))
+            uninterrupted = export(clean)
+            self.assertEqual(resumed['files'], uninterrupted['files'])
+            self.assertEqual(validate(config)['examples_checked'], 150)
+
+    def test_corrupted_shard_is_detected(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            config = Config(output_dir=Path(tmp), target_examples=4)
+            asyncio.run(generate(config))
+            export(config)
+            next(Path(tmp).glob('sft/*/*/*.jsonl')).write_text('broken\n')
+            with self.assertRaisesRegex(ValueError, 'Повреждён'):
+                validate(config)
+
+    def test_all_steps_and_one_mode_obey_total(self):
+        for modes in (['step'], ['branch'], ['full', 'step', 'branch']):
+            with self.subTest(modes=modes), tempfile.TemporaryDirectory() as tmp:
+                config = Config(output_dir=Path(tmp), target_examples=37, modes=modes, step_sampling='all')
+                asyncio.run(generate(config))
+                report = export(config)
+                self.assertEqual(report['examples'], 37)
+                self.assertEqual(validate(config)['examples_checked'], 37)
+
+    def test_smaller_export_does_not_remove_canonical_episodes(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            config = Config(output_dir=Path(tmp), target_examples=100)
+            asyncio.run(generate(config))
+            original = export(config)
+            config.target_examples = 1
+            self.assertEqual(export(config)['examples'], 1)
+            config.target_examples = 100
+            self.assertEqual(export(config)['files'], original['files'])
+
+
+if __name__ == '__main__':
+    unittest.main()