nl_translator.py

"""
Natural Language to MCP Command Translator
Uses Qwen2.5-Coder-1.5B to translate player natural language commands to MCP JSON
Supports English, French, and Chinese
Uses shared model manager to avoid duplicate loading with AI analysis
"""
import json
import re
import time
from typing import Dict, Optional, Tuple
from pathlib import Path

from model_manager import get_shared_model

class NLCommandTranslator:
    """Translates natural language commands to MCP JSON format"""
    
    def __init__(self, model_path: str = "qwen2.5-coder-1.5b-instruct-q4_0.gguf"):
        self.model_path = model_path
        self.model_manager = get_shared_model()
        self.last_error = None
        
        # Language detection patterns
        self.lang_patterns = {
            'zh': re.compile(r'[\u4e00-\u9fff]'),  # Chinese characters
            'fr': re.compile(r'[àâçèéêëîïôùûü]', re.IGNORECASE)  # French accents
        }
        
        # System prompts for each language
        self.system_prompts = {
            "en": """You are an AI assistant for an RTS game. Convert user commands into JSON tool calls.

Available tools:
- get_game_state(): Get current game state
- move_units(unit_ids: list, target_x: int, target_y: int): Move units to position
- attack_unit(attacker_ids: list, target_id: str): Attack enemy unit
- build_unit(unit_type: str): Build a unit (infantry, tank, helicopter, harvester)
- build_building(building_type: str, x: int, y: int): Build a building (barracks, war_factory, power_plant, refinery, defense_turret)

Respond ONLY with valid JSON containing "tool" and "params" fields.
For parameterless functions, you may omit the params field.
Example: {"tool": "move_units", "params": {"unit_ids": ["unit_1"], "target_x": 200, "target_y": 300}}""",
            
            "fr": """Tu es un assistant IA pour un jeu RTS. Convertis les commandes utilisateur en appels d'outils JSON.

Outils disponibles :
- get_game_state(): Obtenir l'état du jeu
- move_units(unit_ids: list, target_x: int, target_y: int): Déplacer des unités
- attack_unit(attacker_ids: list, target_id: str): Attaquer une unité ennemie
- build_unit(unit_type: str): Construire une unité (infantry, tank, helicopter, harvester)
- build_building(building_type: str, x: int, y: int): Construire un bâtiment (barracks, war_factory, power_plant, refinery, defense_turret)

Réponds UNIQUEMENT avec du JSON valide contenant "tool" et "params".
Pour les fonctions sans paramètres, tu peux omettre le champ params.
Exemple: {"tool": "move_units", "params": {"unit_ids": ["unit_1"], "target_x": 200, "target_y": 300}}""",
            
            "zh": """你是一个RTS游戏的AI助手。将用户命令转换为JSON工具调用。

可用工具：
- get_game_state(): 获取游戏状态
- move_units(unit_ids: list, target_x: int, target_y: int): 移动单位
- attack_unit(attacker_ids: list, target_id: str): 攻击敌方单位
- build_unit(unit_type: str): 建造单位 (infantry, tank, helicopter, harvester)
- build_building(building_type: str, x: int, y: int): 建造建筑 (barracks, war_factory, power_plant, refinery, defense_turret)

只返回包含"tool"和"params"字段的有效JSON。
对于无参数函数，可以省略params字段。
示例: {"tool": "move_units", "params": {"unit_ids": ["unit_1"], "target_x": 200, "target_y": 300}}"""
        }
    
    def load_model(self) -> Tuple[bool, Optional[str]]:
        """Load the shared LLM model"""
        return self.model_manager.load_model(self.model_path)
    
    @property
    def model_loaded(self) -> bool:
        """Check if model is loaded"""
        return self.model_manager.model_loaded
    
    def detect_language(self, text: str) -> str:
        """Detect language from text"""
        # Check for Chinese
        if self.lang_patterns['zh'].search(text):
            return 'zh'
        # Check for French
        elif self.lang_patterns['fr'].search(text):
            return 'fr'
        # Default to English
        else:
            return 'en'
    
    def extract_json_from_response(self, text: str) -> Optional[Dict]:
        """Extract JSON from model response"""
        # Try to find JSON in markdown code block
        json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', text, re.DOTALL)
        if json_match:
            try:
                return json.loads(json_match.group(1))
            except json.JSONDecodeError:
                pass
        
        # Try to find raw JSON object
        json_match = re.search(r'\{[^{}]*"tool"[^{}]*\}', text, re.DOTALL)
        if json_match:
            try:
                return json.loads(json_match.group(0))
            except json.JSONDecodeError:
                pass
        
        # Try to parse entire text as JSON
        try:
            return json.loads(text)
        except json.JSONDecodeError:
            return None
    
    def translate_command(self, nl_command: str, language: Optional[str] = None) -> Dict:
        """
        Translate natural language command to MCP JSON
        
        Returns dict with:
        - success: bool
        - json_command: dict (if success)
        - error: str (if not success)
        - language: str (detected language)
        - response_time: float
        """
        start_time = time.time()
        
        # Auto-detect language if not provided
        if language is None:
            language = self.detect_language(nl_command)
        
        # Ensure model is loaded
        if not self.model_loaded:
            success, error = self.load_model()
            if not success:
                return {
                    "success": False,
                    "error": error,
                    "language": language,
                    "response_time": time.time() - start_time
                }
        
        # Get system prompt for language
        system_prompt = self.system_prompts.get(language, self.system_prompts["en"])
        
        # Create chat messages
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": nl_command}
        ]
        
        try:
            # Generate response using shared model
            success, raw_response, error = self.model_manager.generate(
                messages=messages,
                max_tokens=128,  # Reduced for faster response
                temperature=0.1,
                timeout=10.0  # Shorter timeout for game responsiveness
            )
            
            if not success or not raw_response:
                return {
                    "success": False,
                    "error": error or "Model generation failed",
                    "language": language,
                    "response_time": time.time() - start_time
                }
            
            raw_response = raw_response.strip()
            
            # Extract JSON
            json_command = self.extract_json_from_response(raw_response)
            
            response_time = time.time() - start_time
            
            if json_command and 'tool' in json_command:
                return {
                    "success": True,
                    "json_command": json_command,
                    "language": language,
                    "response_time": response_time,
                    "raw_response": raw_response
                }
            else:
                return {
                    "success": False,
                    "error": "Failed to extract valid JSON command",
                    "language": language,
                    "response_time": response_time,
                    "raw_response": raw_response
                }
        
        except Exception as e:
            return {
                "success": False,
                "error": f"Translation error: {str(e)}",
                "language": language,
                "response_time": time.time() - start_time
            }
    
    def get_example_commands(self, language: str = "en") -> list:
        """Get example commands for the given language"""
        examples = {
            "en": [
                "Show me the game state",
                "Move my infantry to position 200, 300",
                "Build a tank",
                "Construct a power plant at 150, 150",
                "Attack the enemy base",
            ],
            "fr": [
                "Montre-moi l'état du jeu",
                "Déplace mon infanterie vers 200, 300",
                "Construis un char",
                "Construit une centrale électrique à 150, 150",
                "Attaque la base ennemie",
            ],
            "zh": [
                "显示游戏状态",
                "移动我的步兵到200, 300",
                "建造一个坦克",
                "在150, 150建造发电厂",
                "攻击敌人的基地",
            ]
        }
        return examples.get(language, examples["en"])


# Global translator instance
_translator = None

def get_nl_translator() -> NLCommandTranslator:
    """Get or create global NL translator instance"""
    global _translator
    if _translator is None:
        _translator = NLCommandTranslator()
        # Ensure model is loaded
        if not _translator.model_loaded:
            print("🔄 Loading NL translator model...")
            success, error = _translator.load_model()
            if success:
                print("✅ NL translator model loaded successfully")
            else:
                print(f"❌ Failed to load NL translator model: {error}")
    return _translator