79 lines
4.1 KiB
Python
79 lines
4.1 KiB
Python
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')
|