prompt_expander_node.py

"""
Main ComfyUI Node for AI Video Prompt Expansion - VERSION 1.2.1
Fixed auto mode + shows detected tier
"""

import os
import re
import random
from typing import Tuple
from .llm_backend import LLMBackend
from .expansion_engine import PromptExpander
from .utils import (
    save_prompts_to_file,
    parse_keywords,
    format_breakdown,
    validate_positive_keywords
)


class AIVideoPromptExpander:
    """
    ComfyUI node that expands simple video ideas into detailed prompts
    using local LLMs (LM Studio or Ollama)
    """
    
    def __init__(self):
        self.expander = PromptExpander()
        self.type = "prompt_expansion"
        self.output_dir = "output/video_prompts"
        self._emphasis_store = []  # Store for emphasis syntax preservation
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                # Core inputs
                "basic_prompt": ("STRING", {
                    "multiline": True,
                    "default": "A cat playing piano in a cozy room"
                }),
                
                "preset": ([
                    "custom",
                    "cinematic", 
                    "surreal",
                    "action",
                    "stylized",
                    "noir",
                    "random"
                ], {
                    "default": "cinematic"
                }),
                
                "expansion_tier": ([
                    "auto",
                    "basic",
                    "enhanced", 
                    "advanced",
                    "cinematic"
                ], {
                    "default": "auto"
                }),
                
                "mode": ([
                    "text-to-video",
                    "image-to-video"
                ], {
                    "default": "text-to-video"
                }),
                
                # LLM Configuration
                "llm_backend": ([
                    "lm_studio",
                    "ollama",
                    "qwen3_vl"
                ], {
                    "default": "lm_studio",
                    "tooltip": (
                        "lm_studio: Uses currently loaded model in LM Studio\n"
                        "ollama: Uses currently loaded model in Ollama\n"
                        "qwen3_vl: Auto-detects local Qwen3-VL model (no API server needed)"
                    )
                }),
                
                "api_endpoint": ("STRING", {
                    "default": "http://localhost:1234/v1",
                    "multiline": False,
                    "tooltip": (
                        "lm_studio/ollama: API endpoint URL\n"
                        "qwen3_vl: Leave default, or specify custom model path like 'A:\\path\\to\\model'"
                    )
                }),
                
                "temperature": ("FLOAT", {
                    "default": 0.7,
                    "min": 0.1,
                    "max": 2.0,
                    "step": 0.1
                }),
                
                # Keywords
                "positive_keywords": ("STRING", {
                    "default": "",
                    "multiline": True,
                    "placeholder": "Enter comma-separated keywords (e.g., lora_trigger, specific_term)"
                }),
                
                "negative_keywords": ("STRING", {
                    "default": "",
                    "multiline": True,
                    "placeholder": "Enter comma-separated negative terms"
                }),
                
                # Output options
                "num_variations": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 3,
                    "step": 1
                }),
                
                "save_to_file": ("BOOLEAN", {
                    "default": False
                }),
                
                "filename_base": ("STRING", {
                    "default": "video_prompt",
                    "multiline": False
                })
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING")
    RETURN_NAMES = (
        "positive_prompt_1",
        "positive_prompt_2", 
        "positive_prompt_3",
        "negative_prompt",
        "breakdown",
        "status"
    )
    
    FUNCTION = "expand_prompt"
    CATEGORY = "Eric Prompt Enhancers"
    OUTPUT_NODE = True
    
    def expand_prompt(
        self,
        basic_prompt: str,
        preset: str,
        expansion_tier: str,
        mode: str,
        llm_backend: str,
        api_endpoint: str,
        temperature: float,
        positive_keywords: str,
        negative_keywords: str,
        num_variations: int,
        save_to_file: bool,
        filename_base: str
    ) -> Tuple[str, str, str, str, str, str]:
        """Main processing function"""
        
        try:
            # Process alternations first (before LLM)
            basic_prompt = self._process_alternations(basic_prompt)
            
            # Preserve emphasis syntax before LLM processing
            basic_prompt = self._preserve_emphasis_syntax(basic_prompt)
            
            pos_kw_list = parse_keywords(positive_keywords)
            neg_kw_list = parse_keywords(negative_keywords)
            
            # Initialize LLM backend (model_name auto-detected)
            llm = LLMBackend(
                backend_type=llm_backend,
                endpoint=api_endpoint,
                model_name=None,  # Auto-detect for all backends
                temperature=temperature
            )
            
            conn_test = llm.test_connection()
            if not conn_test["success"]:
                error_msg = f"LLM Connection Failed: {conn_test['message']}"
                return (basic_prompt, "", "", "", f"ERROR: {error_msg}", f"❌ {error_msg}")
            
            positive_prompts = []
            breakdowns = []
            
            for var_num in range(num_variations):
                system_prompt, user_prompt, breakdown_dict = self.expander.expand_prompt(
                    basic_prompt=basic_prompt,
                    preset=preset,
                    tier=expansion_tier,
                    mode=mode,
                    positive_keywords=pos_kw_list,
                    variation_seed=var_num if num_variations > 1 else None
                )
                
                response = llm.send_prompt(
                    system_prompt=system_prompt,
                    user_prompt=user_prompt,
                    max_tokens=3000
                )
                
                if not response["success"]:
                    error_msg = response["error"]
                    return (basic_prompt, "", "", "", f"ERROR: {error_msg}", f"❌ {error_msg}")
                
                parsed = self.expander.parse_llm_response(response["response"])
                enhanced_prompt = parsed["prompt"]
                
                # Restore emphasis syntax after LLM processing
                enhanced_prompt = self._restore_emphasis_syntax(enhanced_prompt)
                
                # Validate we got output
                if not enhanced_prompt or len(enhanced_prompt) < 20:
                    return (
                        basic_prompt,
                        "",
                        "",
                        "",
                        f"ERROR: LLM returned empty or very short response. Raw: {response['response'][:200]}",
                        f"❌ LLM response too short - check your model"
                    )
                
                if pos_kw_list:
                    keywords_present, missing = validate_positive_keywords(pos_kw_list, enhanced_prompt)
                    if missing:
                        enhanced_prompt += f" {', '.join(missing)}"
                
                positive_prompts.append(enhanced_prompt)
                breakdowns.append(breakdown_dict)
            
            while len(positive_prompts) < 3:
                positive_prompts.append("")
            
            negative_prompt = self.expander.generate_negative_prompt(
                preset=preset,
                custom_negatives=neg_kw_list,
                mode=mode
            )
            
            breakdown_text = self._format_all_breakdowns(breakdowns, basic_prompt)
            
            if save_to_file and positive_prompts[0]:
                metadata = {
                    "preset": preset,
                    "tier": expansion_tier,
                    "detected_tier": breakdowns[0].get('detected_tier'),
                    "mode": mode,
                    "backend": llm_backend,
                    "model": llm.model_name or "auto-detected",
                    "temperature": temperature,
                    "variation_num": num_variations,
                    "original_prompt": basic_prompt
                }
                
                save_result = save_prompts_to_file(
                    positive_prompt=positive_prompts[0],
                    negative_prompt=negative_prompt,
                    breakdown=breakdown_text,
                    metadata=metadata,
                    filename_base=filename_base,
                    output_dir=self.output_dir
                )
                
                file_status = f"💾 Saved: {save_result['filepath']}" if save_result["success"] else f"⚠️ Save failed"
            else:
                file_status = "Not saved"
            
            # Build status - show detected tier if auto was used
            detected_tier = breakdowns[0].get('detected_tier')
            was_auto = breakdowns[0].get('was_auto_detected', False)
            
            if was_auto and detected_tier:
                tier_display = f"Tier: auto→{detected_tier}"
            else:
                tier_display = f"Tier: {expansion_tier}"
            
            status = f"✅ Generated {num_variations} variation(s) | {tier_display} | Preset: {preset} | {file_status}"
            
            return (
                positive_prompts[0],
                positive_prompts[1],
                positive_prompts[2],
                negative_prompt,
                breakdown_text,
                status
            )
        
        except Exception as e:
            import traceback
            error_detail = traceback.format_exc()
            print(f"ERROR in prompt expansion: {error_detail}")
            return (
                basic_prompt,
                "",
                "",
                "",
                f"ERROR: {str(e)}\n\n{error_detail}",
                f"❌ {str(e)}"
            )
    
    def _format_all_breakdowns(self, breakdowns: list, original: str) -> str:
        """Format breakdown information"""
        
        if not breakdowns:
            return "No breakdown available"
        
        lines = [
            "=" * 60,
            "PROMPT EXPANSION BREAKDOWN",
            "=" * 60,
            f"\nOriginal Input:\n{original}\n",
        ]
        
        # Show if auto-detected
        if breakdowns[0].get('was_auto_detected'):
            lines.append(f"Tier Setting: auto (detected: {breakdowns[0].get('detected_tier')})")
        else:
            lines.append(f"Detected Tier: {breakdowns[0].get('detected_tier', 'N/A')}")
        
        lines.extend([
            f"Applied Preset: {breakdowns[0].get('applied_preset', 'N/A')}",
            f"Mode: {breakdowns[0].get('mode', 'N/A')}",
        ])
        
        if breakdowns[0].get('positive_keywords'):
            lines.append(f"Required Keywords: {', '.join(breakdowns[0]['positive_keywords'])}")
        
        if breakdowns[0].get('preset_focus'):
            lines.append(f"Focus Areas: {', '.join(breakdowns[0]['preset_focus'])}")
        
        lines.append("\n" + "=" * 60)
        
        return "\n".join(lines)
    
    def _process_alternations(self, text: str) -> str:
        """
        Process alternation syntax {option1|option2|option3}
        Replaces with randomly chosen option
        """
        # Pattern to match {option1|option2|option3}
        pattern = r'\{([^{}]+)\}'
        
        def replace_alternation(match):
            options = match.group(1).split('|')
            # Strip whitespace from each option
            options = [opt.strip() for opt in options]
            return random.choice(options)
        
        # Keep replacing until no more alternations found (handles nested cases)
        max_iterations = 10  # Prevent infinite loops
        iteration = 0
        while '{' in text and '|' in text and iteration < max_iterations:
            new_text = re.sub(pattern, replace_alternation, text)
            if new_text == text:  # No more changes
                break
            text = new_text
            iteration += 1
        
        return text
    
    def _preserve_emphasis_syntax(self, text: str) -> str:
        """
        Protect emphasis syntax (keyword:1.5) from being modified
        Replaces temporarily with placeholders during LLM processing
        """
        # Pattern to match (text:number) emphasis syntax
        # This matches things like (dark skin:1.5) or (hair:0.8)
        pattern = r'\(([^():]+):(\d+\.?\d*)\)'
        
        # Find all emphasis patterns
        emphasis_patterns = re.findall(pattern, text)
        
        # Store original patterns
        self._emphasis_store = []
        
        # Replace with placeholders
        def replace_emphasis(match):
            full_match = match.group(0)
            placeholder = f"__EMPHASIS_{len(self._emphasis_store)}__"
            self._emphasis_store.append(full_match)
            return placeholder
        
        text = re.sub(pattern, replace_emphasis, text)
        
        return text
    
    def _restore_emphasis_syntax(self, text: str) -> str:
        """
        Restore emphasis syntax that was protected
        """
        if not hasattr(self, '_emphasis_store'):
            return text
        
        # Restore placeholders with original emphasis syntax
        for i, original in enumerate(self._emphasis_store):
            placeholder = f"__EMPHASIS_{i}__"
            text = text.replace(placeholder, original)
        
        # Clear the store
        self._emphasis_store = []
        
        return text