ai_qc/backend/profile_config.py

#!/usr/bin/env python3
"""
Centralized profile configuration module for Visual AI QC.
This script manages the profiles, QC checks, weights, and LLM assignments.
"""

import os
import json
import glob
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field

# Profiles directory path
PROFILES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'profiles')

# Dynamic QC checks discovery
def discover_qc_checks():
    """Dynamically discover all available QC checks from the visual_qc_apps directory"""
    import glob
    
    visual_qc_apps_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'visual_qc_apps')
    
    # Find all directories that contain an app.py file
    qc_checks = []
    
    # Look for all subdirectories in visual_qc_apps
    app_dirs = glob.glob(os.path.join(visual_qc_apps_dir, '*', 'app.py'))
    
    for app_file in app_dirs:
        # Extract the directory name as the check name
        check_dir = os.path.dirname(app_file)
        check_name = os.path.basename(check_dir)
        
        # Skip template files and utility files
        if check_name not in ['__pycache__', 'templates', 'utils']:
            qc_checks.append(check_name)
    
    # Sort for consistency
    qc_checks.sort()
    
    print(f"Discovered {len(qc_checks)} QC checks: {qc_checks}")
    return qc_checks

# List of all available QC checks (dynamically discovered)
try:
    QC_CHECKS = discover_qc_checks()
except Exception as e:
    print(f"Error discovering QC checks, falling back to static list: {e}")
    # Fallback to static list if discovery fails
    QC_CHECKS = [
        # Original services
        'logo_visibility',
        'brand_assets_visibility',
        'visual_elements_count',
        'background_contrast',
        'face_visibility',
        'new_visibility',
        'visual_hierarchy',
        'supporting_images',
        'curved_edges',
        'visuals_left_text_right',
        'face_gaze_direction',
        'lowercase_text',
        'call_to_action',
        'word_count',
        'imperative_verb',
        
        # New services - added for the new client requirements
        'file_naming',
        'layer_organization',
        'color_format',
        'image_resolution',
        'safety_area',
        'element_alignment',
        'animation_transitions',
        'aspect_ratio',
        'responsiveness',
        'dark_mode_legibility',
        'print_bleed',
        'crop_marks',
        'text_readability',
        
        # Custom QC checks for specific requirements
        'product_visibility',
        'inclusive',
        'accessibility',
        
        # Format-specific checks
        'curved_edges_print',
        'curved_edges_digital'
    ]

# LLM options
LLM_OPTIONS = ["OpenAI", "Gemini"]

@dataclass
class QCCheckConfig:
    """Configuration for a QC check, including its weight and which LLM to use"""
    weight: float = 0.0
    llm: str = "Gemini"  # Default to Gemini
    enabled: bool = True

@dataclass
class Profile:
    """Profile configuration including name, description, and check configs"""
    name: str
    description: str
    checks: Dict[str, QCCheckConfig] = field(default_factory=dict)
    pre_analysis_instructions: Optional[str] = None

    def get_enabled_checks(self) -> List[str]:
        """Get list of enabled check names"""
        return [check_name for check_name, config in self.checks.items() if config.enabled]
    
    def get_check_weights(self) -> Dict[str, float]:
        """Get dictionary of check weights"""
        return {check_name: config.weight for check_name, config in self.checks.items() if config.enabled}
    
    def get_check_llm(self, check_name: str) -> str:
        """Get the LLM to use for a specific check"""
        if check_name in self.checks:
            return self.checks[check_name].llm
        return "Gemini"  # Default to Gemini if not specified

# Dictionary to store all loaded profiles
PROFILES = {}

def load_profiles():
    """Load all profile JSON files from the profiles directory"""
    global PROFILES
    PROFILES = {}  # Reset profiles dictionary
    
    # Ensure profiles directory exists
    os.makedirs(PROFILES_DIR, exist_ok=True)
    
    # Find all JSON files in the profiles directory
    profile_files = glob.glob(os.path.join(PROFILES_DIR, '*.json'))
    
    # Load each profile file
    for profile_file in profile_files:
        try:
            with open(profile_file, 'r') as f:
                profile_data = json.load(f)
                
            # Extract profile name, description, checks, and pre_analysis_instructions
            profile_name = profile_data.get('name', 'Unnamed Profile')
            profile_description = profile_data.get('description', '')
            profile_checks = profile_data.get('checks', {})
            pre_analysis_instructions = profile_data.get('pre_analysis_instructions', None)
            
            # Create a new Profile instance
            profile = Profile(
                name=profile_name,
                description=profile_description,
                pre_analysis_instructions=pre_analysis_instructions
            )
            
            # Add each check configuration
            for check_name, check_config in profile_checks.items():
                profile.checks[check_name] = QCCheckConfig(
                    weight=check_config.get('weight', 0.0),
                    llm=check_config.get('llm', 'Gemini'),
                    enabled=check_config.get('enabled', True)
                )
            
            # Add profile to the PROFILES dictionary
            # Use the filename (without extension) as the profile ID
            profile_id = os.path.splitext(os.path.basename(profile_file))[0].lower()
            PROFILES[profile_id] = profile
            print(f"Loaded profile '{profile_name}' from {profile_file}")
            
        except Exception as e:
            print(f"Error loading profile from {profile_file}: {e}")
    
    # If no profiles were loaded, create a default profile
    if not PROFILES:
        print("No profiles found. Creating default profile.")
        default_profile = Profile(
            name="All Checks",
            description="Run all available QC checks"
        )
        # Initialize all checks with default values
        for check in QC_CHECKS:
            default_profile.checks[check] = QCCheckConfig()
        
        PROFILES['default'] = default_profile
        
        # Save the default profile to a file
        save_profile('default', default_profile)

def save_profile(profile_id: str, profile: Profile):
    """Save a profile to a JSON file"""
    # Create the profile data dictionary
    profile_data = {
        'name': profile.name,
        'description': profile.description,
        'checks': {}
    }
    
    # Add pre_analysis_instructions if it exists
    if profile.pre_analysis_instructions:
        profile_data['pre_analysis_instructions'] = profile.pre_analysis_instructions
    
    # Add each check configuration
    for check_name, check_config in profile.checks.items():
        profile_data['checks'][check_name] = {
            'weight': check_config.weight,
            'llm': check_config.llm,
            'enabled': check_config.enabled
        }
    
    # Save to a JSON file
    profile_file = os.path.join(PROFILES_DIR, f"{profile_id.lower()}.json")
    with open(profile_file, 'w') as f:
        json.dump(profile_data, f, indent=4)
    
    print(f"Saved profile '{profile.name}' to {profile_file}")

def get_profile(profile_name: str) -> Profile:
    """Get a profile by name"""
    # If profiles haven't been loaded yet, load them
    if not PROFILES:
        load_profiles()
        
    return PROFILES.get(profile_name.lower(), PROFILES.get('default'))

def add_profile(name: str, description: str, check_configs: Dict[str, Dict[str, Any]]) -> str:
    """Add a new profile and save it to a JSON file
    
    Returns the profile_id that was created
    """
    # Create a new Profile instance
    profile = Profile(
        name=name,
        description=description
    )
    
    # Add each check configuration
    for check_name, config in check_configs.items():
        profile.checks[check_name] = QCCheckConfig(
            weight=config.get('weight', 0.0),
            llm=config.get('llm', 'Gemini'),
            enabled=config.get('enabled', True)
        )
    
    # Generate a profile_id from the name
    profile_id = name.lower().replace(' ', '_')
    
    # Add to the PROFILES dictionary
    PROFILES[profile_id] = profile
    
    # Save to a JSON file
    save_profile(profile_id, profile)
    
    return profile_id

def update_profile(profile_name: str, updates: Dict[str, Any]) -> bool:
    """Update an existing profile"""
    if profile_name not in PROFILES:
        return False
    
    profile = PROFILES[profile_name]
    
    if 'name' in updates:
        profile.name = updates['name']
    
    if 'description' in updates:
        profile.description = updates['description']
    
    if 'checks' in updates:
        for check_name, check_config in updates['checks'].items():
            if check_name in profile.checks:
                if 'weight' in check_config:
                    profile.checks[check_name].weight = check_config['weight']
                if 'llm' in check_config:
                    profile.checks[check_name].llm = check_config['llm']
                if 'enabled' in check_config:
                    profile.checks[check_name].enabled = check_config['enabled']
    
    # Save the updated profile
    save_profile(profile_name, profile)
    
    return True

def delete_profile(profile_name: str) -> bool:
    """Delete a profile file and remove it from memory"""
    if profile_name not in PROFILES or profile_name == 'default':
        return False
    
    # Remove from memory
    profile = PROFILES.pop(profile_name)
    
    # Delete the file
    profile_file = os.path.join(PROFILES_DIR, f"{profile_name.lower()}.json")
    if os.path.exists(profile_file):
        os.remove(profile_file)
        print(f"Deleted profile file: {profile_file}")
    
    return True

def get_profile_summary() -> Dict[str, Dict[str, Any]]:
    """Get a summary of all profiles"""
    # If profiles haven't been loaded yet, load them
    if not PROFILES:
        load_profiles()
        
    summary = {}
    
    for profile_name, profile in PROFILES.items():
        summary[profile_name] = {
            'name': profile.name,
            'description': profile.description,
            'enabled_checks': profile.get_enabled_checks(),
            'total_checks': len(profile.checks),
            'enabled_count': len(profile.get_enabled_checks())
        }
    
    return summary

def get_check_llm_map(profile_name: str) -> Dict[str, str]:
    """Get a mapping of check names to LLM names for a profile"""
    profile = get_profile(profile_name)
    return {check_name: config.llm for check_name, config in profile.checks.items()}

# Load profiles when the module is imported
load_profiles()