78 lines
4.3 KiB
Python
78 lines
4.3 KiB
Python
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})
|