72 lines
3.4 KiB
Python
72 lines
3.4 KiB
Python
"""Optional teacher plans through the project's existing model factory."""
|
||
import asyncio
|
||
import re
|
||
from typing import Literal
|
||
|
||
from pydantic import Field, field_validator
|
||
from llm.config import ConfigModel
|
||
from .prompts import render
|
||
from .language import validate_texts, action_text
|
||
from .world import canonical, validate_plan, branch_plan, validate_branch
|
||
|
||
|
||
class Action(ConfigModel):
|
||
skill: str
|
||
object: str
|
||
target: str
|
||
text: str = Field(description='Одно действие на русском языке с названием объекта из сцены; без кода')
|
||
|
||
@field_validator('text')
|
||
@classmethod
|
||
def russian_description(cls, value):
|
||
value = value.strip()
|
||
if not re.search('[А-Яа-яЁё]', value) or re.search(r'\b(?:part_|slot_)\d+|\b(?:robot|camera)\.', value):
|
||
raise ValueError('Ожидается русское описание действия, без технических идентификаторов')
|
||
return value
|
||
|
||
|
||
class Plan(ConfigModel):
|
||
status: Literal['ready', 'blocked']
|
||
actions: list[Action] = Field(default_factory=list)
|
||
|
||
|
||
async def create_agent(config, prompts):
|
||
from pydantic_ai import Agent
|
||
from llm.factory import LLMModelProvider
|
||
model = await LLMModelProvider.create_from_config(config.llm)
|
||
return Agent(model=model, output_type=Plan, retries=config.retries,
|
||
system_prompt=prompts.teacher_system)
|
||
|
||
|
||
async def enrich(episode, config, agent, prompts):
|
||
if int(episode['signature'][16:32], 16) / 2**64 >= config.fraction:
|
||
return episode, None
|
||
teacher_scene = {**episode['scene'], 'operations': [
|
||
{**operation, 'reference_text': action_text(operation['action'], episode['scene'], prompts)}
|
||
for operation in episode['scene']['operations']]}
|
||
prompt = render(prompts.teacher_user, canonical({'instruction': episode['instruction'], 'scene': teacher_scene}))
|
||
error = None
|
||
for attempt in range(config.retries + 1):
|
||
try:
|
||
async with asyncio.timeout(config.timeout_seconds):
|
||
result = await agent.run(prompt, model_settings=config.generation.model_dump(exclude_none=True))
|
||
raw = result.output.model_dump()
|
||
texts = [action['text'] for action in raw['actions']]
|
||
plan = {'status': raw['status'], 'actions': [{k: v for k, v in a.items() if k != 'text'} for a in raw['actions']]}
|
||
validate_texts({**episode, 'plan': plan}, texts, prompts)
|
||
trace = validate_plan(episode, plan)
|
||
branch = branch_plan(plan)
|
||
validate_branch(episode, branch)
|
||
episode.update(plan=plan, trace=trace, branch=branch, step_texts=texts, text_source='teacher',
|
||
source={'provider': config.llm.provider, 'model': config.llm.model})
|
||
return episode, None
|
||
except Exception as exc:
|
||
# API exceptions may contain request headers. Persist the class only.
|
||
error = {'scenario_id': episode['scenario_id'], 'error_type': type(exc).__name__, 'attempts': attempt + 1}
|
||
if isinstance(exc, ValueError):
|
||
prompt = render(prompts.teacher_retry, canonical({'instruction': episode['instruction'], 'scene': teacher_scene,
|
||
'validation_error': str(exc)[:500]}))
|
||
# A failed teacher is never silently labelled as a successful teacher plan.
|
||
episode['teacher_error'] = error
|
||
return episode, error
|