gensim/code_templates.py

"""
Sistema de templates flexíveis para estruturas comuns.
Templates definem estrutura básica, LLM preenche lógica específica.
"""
from typing import Dict, Optional, List
import re


class CodeTemplate:
    """Template para estruturas comuns de código"""
    
    TEMPLATES = {
        'row': {
            'description': 'Template para criar fileira de blocos',
            'structure': """
# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'stacking/block.urdf'
block_color = utils.COLORS['{color}']

objs = []
for _ in range({count}):
    block_pose = self.get_random_pose(env, block_size)
    block_id = env.add_object(block_urdf, block_pose, color=block_color)
    objs.append(block_id)

# Define target positions for the blocks in a row.
base_pose = ((0.5, 0.0, 0.02), (0, 0, 0, 1))
offsets = {offsets_placeholder}  # LLM generates: e.g., [(-0.05, 0, 0.02), (0, 0, 0.02), (0.05, 0, 0.02)]
targs = [(utils.apply(base_pose, offset), base_pose[1]) for offset in offsets]

# Goal: blocks are arranged in a row.
self.add_goal(objs=objs, matches=np.eye({count}), targ_poses=targs, replace=False,
              rotations=True, metric='pose', params=None, step_max_reward=1.0, language_goal=self.lang_template)
"""
        },
        'stack': {
            'description': 'Template para empilhar blocos',
            'structure': """
# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'stacking/block.urdf'
block_color = utils.COLORS['{color}']

objs = []
for _ in range({count}):
    block_pose = self.get_random_pose(env, block_size)
    block_id = env.add_object(block_urdf, block_pose, color=block_color)
    objs.append(block_id)

# Define target positions for stacking.
base_pose = ((0.5, 0.0, 0.02), (0, 0, 0, 1))
stack_poses = {stack_poses_placeholder}  # LLM generates: e.g., [(0, 0, 0.04), (0, 0, 0.08), (0, 0, 0.12)]
targs = [(utils.apply(base_pose, pos), base_pose[1]) for pos in stack_poses]

# Goal: blocks are stacked.
self.add_goal(objs=objs, matches=np.eye({count}), targ_poses=targs, replace=False,
              rotations=True, metric='pose', params=None, step_max_reward=1.0, language_goal=self.lang_template)
"""
        },
        'pyramid': {
            'description': 'Template para pirâmide de blocos',
            'structure': """
# Add base.
base_size = (0.05, 0.15, 0.005)
base_urdf = 'stacking/stand.urdf'
base_pose = self.get_random_pose(env, base_size)
env.add_object(base_urdf, base_pose, category='fixed')

# Add blocks.
block_size = (0.04, 0.04, 0.04)
block_urdf = 'stacking/block.urdf'
block_color = utils.COLORS['{color}']

objs = []
for _ in range({count}):
    block_pose = self.get_random_pose(env, block_size)
    block_id = env.add_object(block_urdf, block_pose, color=block_color)
    objs.append(block_id)

# IMPORTANT: Associate placement locations for goals.
# Pyramid structure: {pyramid_structure_placeholder}  # LLM describes structure
place_pos = {place_pos_placeholder}  # LLM generates: e.g., [(0, -0.05, 0.03), (0, 0, 0.03), (0, 0.05, 0.03), ...]
targs = [(utils.apply(base_pose, i), base_pose[1]) for i in place_pos]

# Goal: blocks are stacked in a pyramid.
# Use multiple add_goal() calls for each layer
{goals_placeholder}  # LLM generates multiple add_goal() calls
"""
        }
    }
    
    @classmethod
    def detect_template_type(cls, task_description: str, task_name: str) -> Optional[str]:
        """
        Detecta qual template usar baseado na descrição da tarefa
        
        Returns:
            Nome do template ou None se não detectar
        """
        desc_lower = task_description.lower()
        name_lower = task_name.lower()
        
        # Detectar fileira/linha
        if any(word in desc_lower or word in name_lower 
               for word in ['row', 'fileira', 'linha', 'line', 'fila']):
            return 'row'
        
        # Detectar empilhamento simples
        if any(word in desc_lower or word in name_lower 
               for word in ['stack', 'empilhar', 'stacking']):
            if 'pyramid' not in desc_lower and 'pyramid' not in name_lower:
                return 'stack'
        
        # Detectar pirâmide
        if 'pyramid' in desc_lower or 'pyramid' in name_lower or 'pirâmide' in desc_lower:
            return 'pyramid'
        
        return None
    
    @classmethod
    def get_template(cls, template_type: str) -> Optional[Dict]:
        """Retorna template específico"""
        return cls.TEMPLATES.get(template_type)
    
    @classmethod
    def should_use_template(cls, task_description: str, task_name: str) -> bool:
        """Decide se deve usar template ou geração livre"""
        template_type = cls.detect_template_type(task_description, task_name)
        return template_type is not None


def suggest_template_usage(task_description: str, task_name: str) -> Optional[str]:
    """
    Sugere se deve usar template e qual
    
    Returns:
        Nome do template sugerido ou None
    """
    return CodeTemplate.detect_template_type(task_description, task_name)