core/config_generator.py

"""
Configuration generator
Use AI to generate layout configuration from natural language requirements
"""
import yaml
from openai import OpenAI
from pathlib import Path
from typing import Dict, Any

from config.settings import settings
from config.prompts import SYSTEM_PROMPT, get_config_generation_prompt
from utils.logger import logger


class ConfigGenerator:
    """Generate YAML configuration using AI"""
    
    def __init__(self):
        self.client = OpenAI(
            api_key=settings.OPENAI_API_KEY,
            base_url=settings.OPENAI_BASE_URL
        )
    
    def generate(self, user_input: str) -> str:
        """
        Generate YAML configuration from user input
        
        Args:
            user_input: User's natural language requirement description
            
        Returns:
            Generated YAML configuration content (string)
            
        Raises:
            ValueError: If generated YAML format is invalid
        """
        logger.info("Generating configuration from user input...")
        
        prompt = get_config_generation_prompt(user_input)
        
        try:
            response = self.client.chat.completions.create(
                model=settings.OPENAI_MODEL,
                messages=[
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=300,
                temperature=0.8
            )
            
            yaml_content = response.choices[0].message.content.strip()
            
            # Debug: Show what AI generated
            logger.debug(f"AI generated config:\n{yaml_content}")
            
            # Validate generated YAML
            self._validate_yaml(yaml_content)
            
            logger.info("Configuration generated successfully")
            return yaml_content
            
        except Exception as e:
            logger.error(f"Failed to generate configuration: {e}")
            raise
    
    def _validate_yaml(self, yaml_content: str) -> Dict[str, Any]:
        """Validate YAML format"""
        try:
            config = yaml.safe_load(yaml_content)
            if config is None:
                raise ValueError("YAML content is empty")
            return config
        except yaml.YAMLError as e:
            raise ValueError(f"Invalid YAML format: {e}")
    
    def save(self, yaml_content: str, output_path: Path) -> None:
        """Save configuration to file"""
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(yaml_content)
        logger.info(f"Configuration saved to {output_path}")