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

tests/test_prompts.py

@@ -0,0 +1,78 @@
+import asyncio
+import json
+from pathlib import Path
+import tempfile
+from types import SimpleNamespace
+import unittest
+from unittest.mock import AsyncMock, patch
+import yaml
+
+from assembly_dataset.config import Config, load_config
+from assembly_dataset.pipeline import generate
+from assembly_dataset.world import make_episode, training_rows
+from assembly_dataset.language import action_text
+
+
+class PromptTests(unittest.TestCase):
+    def test_custom_request_and_retry_reach_teacher(self):
+        from assembly_dataset.teacher import enrich, Plan
+        from llm.config import OpenAiConfig
+        config = Config(blocked_fraction=0)
+        config.teacher.fraction = 1.0
+        config.teacher.llm = OpenAiConfig(provider='openai', model='test')
+        config.prompts.teacher_user = 'REQUEST ${payload}'
+        config.prompts.teacher_retry = 'REPAIR ${payload}'
+        episode = make_episode(0, config)
+        agent = SimpleNamespace(run=AsyncMock(side_effect=[
+            SimpleNamespace(output=Plan(status='ready', actions=[])),
+            SimpleNamespace(output=Plan.model_validate({**episode['plan'], 'actions': [{**a, 'text': action_text(a, episode['scene'], config.prompts)} for a in episode['plan']['actions']]})),
+        ]))
+        result, error = asyncio.run(enrich(episode, config.teacher, agent, config.prompts))
+        self.assertIsNone(error)
+        first, second = [call.args[0] for call in agent.run.call_args_list]
+        self.assertTrue(first.startswith('REQUEST {'))
+        self.assertTrue(second.startswith('REPAIR {'))
+        self.assertIn('validation_error', json.loads(second.removeprefix('REPAIR ')))
+        self.assertEqual(result['source']['provider'], 'openai')
+
+    def test_external_prompts_relative_paths_and_export(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            prompts = yaml.safe_load(Path('prompts.yaml').read_text())
+            prompts['training_system']['full'] = 'Custom training rules'
+            prompts['training_user'] = 'Input: ${payload}'
+            prompts['task_instructions']['mechanical'] = 'Custom assembly task'
+            (root / 'texts.yaml').write_text(yaml.safe_dump(prompts))
+            (root / 'config.yaml').write_text('prompts_file: texts.yaml\noutput_dir: data\ndomains: [mechanical]\n')
+            config = load_config(root / 'config.yaml')
+            episode = make_episode(0, config)
+            record = next(training_rows(episode, config))
+            self.assertEqual(episode['instruction'], 'Custom assembly task')
+            self.assertEqual(record['messages'][0]['content'], 'Custom training rules')
+            self.assertTrue(record['messages'][1]['content'].startswith('Input: {'))
+            self.assertEqual(config.output_dir, root / 'data')
+
+    def test_prompt_changes_prevent_mixed_resume(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            config = Config(output_dir=Path(tmp), target_examples=1)
+            asyncio.run(generate(config))
+            config.prompts.teacher_system = 'Different rules'
+            with self.assertRaises(ValueError):
+                asyncio.run(generate(config))
+
+    def test_invalid_templates_are_rejected(self):
+        from assembly_dataset.prompts import Prompts
+        data = Config().prompts.model_dump()
+        for template in ('No scene', '${unknown}', '${payload.bad}'):
+            with self.subTest(template=template), self.assertRaises(ValueError):
+                Prompts.model_validate({**data, 'teacher_user': template})
+
+    def test_teacher_system_is_passed_to_agent(self):
+        from assembly_dataset.teacher import create_agent
+        from llm.config import OpenAiConfig
+        config = Config()
+        config.teacher.llm = OpenAiConfig(provider='openai', model='test')
+        config.prompts.teacher_system = 'Custom teacher instructions'
+        with patch('llm.factory.LLMModelProvider.create_from_config', new_callable=AsyncMock), patch('pydantic_ai.Agent') as agent:
+            asyncio.run(create_agent(config.teacher, config.prompts))
+            self.assertEqual(agent.call_args.kwargs['system_prompt'], 'Custom teacher instructions')