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

tests/test_russian_plans.py

@@ -0,0 +1,77 @@
+import asyncio
+import re
+from types import SimpleNamespace
+import unittest
+
+from assembly_dataset.config import Config
+from assembly_dataset.world import make_episode, training_rows, validate_plan
+
+
+class RussianPlanTests(unittest.TestCase):
+    def test_training_answers_are_russian_text_for_all_modes(self):
+        config = Config(blocked_fraction=0)
+        episode = make_episode(0, config)
+        for row in training_rows(episode, config):
+            text = row['messages'][-1]['content']
+            self.assertRegex(text, '[А-Яа-я]')
+            self.assertNotRegex(text, r'part_\d|slot_\d|"skill"|on_success')
+            self.assertFalse(text.startswith('{'))
+            if row['mode'] == 'full':
+                self.assertTrue(text.startswith('1. '))
+                self.assertIn('Захватить', text)
+
+    def test_operator_assistance_holds_until_completion_then_returns_to_storage(self):
+        config = Config(domains=['operator_assistance'], min_components=1, max_components=1, blocked_fraction=0)
+        episode = make_episode(0, config)
+        skills = [a['skill'] for a in episode['plan']['actions']]
+        self.assertEqual(skills, ['detect', 'approach', 'pick', 'move_to_operator', 'hold',
+                                  'wait_assembly', 'move_to_storage', 'release', 'return_home'])
+        validate_plan(episode, episode['plan'])
+        text = next(training_rows(episode, config))['messages'][-1]['content']
+        self.assertIn('распределительная коробка', text)
+        self.assertIn('завершения сборки', text)
+        self.assertIn('хранения', text)
+        self.assertIn('исходное положение', text)
+
+    def test_teacher_wording_is_exported_without_replacing_it(self):
+        from assembly_dataset.teacher import Plan, enrich
+        from assembly_dataset.language import action_text
+        from llm.config import OpenAiConfig
+        config = Config(blocked_fraction=0)
+        config.teacher.fraction = 1
+        config.teacher.llm = OpenAiConfig(provider='openai', model='test')
+        episode = make_episode(0, config)
+        actions = [{**a, 'text': action_text(a, episode['scene'], config.prompts)} for a in episode['plan']['actions']]
+        actions[0]['text'] = 'Сначала ' + actions[0]['text'].lower()
+        output = Plan(status='ready', actions=actions)
+        class Agent:
+            async def run(self, *args, **kwargs):
+                return SimpleNamespace(output=output)
+        enriched, error = asyncio.run(enrich(episode, config.teacher, Agent(), config.prompts))
+        self.assertIsNone(error)
+        text = next(training_rows(enriched, config))['messages'][-1]['content']
+        self.assertIn(actions[0]['text'], text)
+
+    def test_teacher_cannot_supply_only_english_skill_text(self):
+        from assembly_dataset.teacher import Plan
+        with self.assertRaises(ValueError):
+            Plan(status='ready', actions=[{'skill':'pick','object':'part_0','target':'slot_0','text':'pick: part_0 -> slot_0'}])
+
+    def test_russian_text_cannot_change_the_validated_action(self):
+        from assembly_dataset.language import action_text, validate_texts
+        config = Config(domains=['operator_assistance'], min_components=1, max_components=1, blocked_fraction=0)
+        episode = make_episode(0, config)
+        texts = [action_text(a, episode['scene'], config.prompts) for a in episode['plan']['actions']]
+        for bad in ('Бросить объект «распределительная коробка №1» на пол.',
+                    texts[0] + ' Затем отпустить объект.'):
+            with self.subTest(text=bad), self.assertRaises(ValueError):
+                validate_texts(episode, [bad, *texts[1:]], config.prompts)
+
+    def test_approaching_another_object_cannot_interrupt_holding(self):
+        config = Config(domains=['operator_assistance'], min_components=2, max_components=2, blocked_fraction=0)
+        episode = make_episode(0, config)
+        actions = episode['plan']['actions']
+        # Execute the second object's approach while the first is being held.
+        reordered = actions[:5] + actions[9:11] + actions[5:9] + actions[11:]
+        with self.assertRaises(ValueError):
+            validate_plan(episode, {'status': 'ready', 'actions': reordered})