backend/tasks/createBenchConfigFile.py

"""
Task to create and save the configuration file
"""
import os
import pathlib
import uuid
import yaml
import shutil
import time
import threading
from typing import Optional, Dict, Any, List, Tuple

from loguru import logger
from huggingface_hub import HfApi


class CreateBenchConfigTask:
    """
    Task to create and save a configuration file for YourbenchSimpleDemo
    """

    def __init__(self, session_uid: Optional[str] = None):
        """
        Initialize the task with a session ID
        
        Args:
            session_uid: Optional session ID, will be generated if None
        """
        self.session_uid = session_uid or str(uuid.uuid4())
        self.logs: List[str] = []
        self.is_completed = False
        self.is_running_flag = threading.Event()
        self.thread = None
        self._add_log("[INFO] Initializing configuration creation task")
    
    def _add_log(self, message: str) -> None:
        """
        Add a log message to the logs list
        
        Args:
            message: Log message to add
        """
        if message not in self.logs:  # Avoid duplicates
            self.logs.append(message)
            # Force a copy of the list to avoid reference issues
            self.logs = self.logs.copy()
            # Log to system logs
            logger.info(f"[{self.session_uid}] {message}")
    
    def get_logs(self) -> List[str]:
        """
        Get all logs for this task
        
        Returns:
            List of log messages
        """
        return self.logs.copy()  # Retourner une copie pour éviter les problèmes de référence
    
    def save_uploaded_file(self, file_path: str) -> str:
        """
        Process the uploaded file that is already in the correct directory
        
        Args:
            file_path: Path to the uploaded file
            
        Returns:
            Path to the file (same as input)
        """
        try:
            # The file is already in the correct location: uploaded_files/{session_id}/uploaded_files/
            # Just log that we're processing it and return the path
            self._add_log(f"[INFO] Processing file: {os.path.basename(file_path)}")
            return file_path
        except Exception as e:
            error_msg = f"Error processing file: {str(e)}"
            self._add_log(f"[ERROR] {error_msg}")
            raise RuntimeError(error_msg)
    
    def generate_base_config(self, hf_org: str, hf_dataset_name: str) -> Dict[str, Any]:
        """
        Create the base configuration dictionary
        
        Args:
            hf_org: Hugging Face organization name
            hf_dataset_name: Hugging Face dataset name
            
        Returns:
            Configuration dictionary
        """
        self._add_log(f"[INFO] Generating base configuration for {hf_dataset_name}")
        
        # Check if HF token is available
        hf_token = os.getenv("HF_TOKEN")
        if not hf_token:
            raise RuntimeError("HF_TOKEN environment variable is not defined")
        
        return {
            "hf_configuration": {
                "token": "$HF_TOKEN",  # Utiliser directement le token de l'environnement
                "hf_organization": "$HF_ORGANIZATION",
                "private": True,
                "hf_dataset_name": hf_dataset_name,
                "concat_if_exist": False,
            },
            "model_list": [
                {
                    "model_name": "Qwen/Qwen2.5-VL-72B-Instruct",
                    "provider": "novita",
                    "api_key": "$HF_TOKEN",
                    "max_concurrent_requests": 32,
                },
                {
                    "model_name": "Qwen/Qwen2.5-72B-Instruct",
                    "provider": "novita",
                    "api_key": "$HF_TOKEN",
                    "max_concurrent_requests": 32,
                },
            ],
            
            "model_roles": {
                "ingestion": ["Qwen/Qwen2.5-VL-72B-Instruct"],
                "summarization": ["Qwen/Qwen2.5-72B-Instruct"],
                "chunking": ["intfloat/multilingual-e5-large-instruct"],
                "single_shot_question_generation": ["Qwen/Qwen2.5-72B-Instruct"],
                "multi_hop_question_generation": ["Qwen/Qwen2.5-72B-Instruct"],
            },
            "pipeline": {
                "ingestion": {
                    "source_documents_dir": f"uploaded_files/{self.session_uid}/uploaded_files/",
                    "output_dir": f"uploaded_files/{self.session_uid}/ingested",
                    "run": True,
                },
                "upload_ingest_to_hub": {
                    "source_documents_dir": f"uploaded_files/{self.session_uid}/ingested",
                    "run": True,  # Réactivé pour l'upload sur le Hub
                },
                "summarization": {
                    "run": True,
                },
                "chunking": {
                    "run": True,
                    "chunking_configuration": {
                        "l_min_tokens": 64,
                        "l_max_tokens": 128,
                        "tau_threshold": 0.8,
                        "h_min": 2,
                        "h_max": 5,
                        "num_multihops_factor": 2,
                    },
                },
                "single_shot_question_generation": {
                    "run": True,
                    "additional_instructions": "Generate questions to test a curious adult",
                    "chunk_sampling": {
                        "mode": "count",
                        "value": 5,
                        "random_seed": 123,
                    },
                },
                "multi_hop_question_generation": {
                    "run": True,
                    "additional_instructions": "Generate questions to test a curious adult",
                    "chunk_sampling": {
                        "mode": "percentage",
                        "value": 0.3,
                        "random_seed": 42,
                    },
                },
                "lighteval": {
                    "run": True,
                },
            },
        }
    
    def save_yaml_file(self, config: Dict[str, Any], path: str) -> str:
        """
        Save the given configuration dictionary to a YAML file
        
        Args:
            config: Configuration dictionary
            path: Path to save the file
            
        Returns:
            Path to the saved file
        """
        try:
            # Create directory if it doesn't exist
            os.makedirs(os.path.dirname(path), exist_ok=True)
            
            with open(path, "w") as file:
                yaml.dump(config, file, default_flow_style=False, sort_keys=False)
            
            self._add_log(f"[INFO] Configuration saved: {path}")
            return path
        except Exception as e:
            error_msg = f"Error saving configuration: {str(e)}"
            self._add_log(f"[ERROR] {error_msg}")
            raise RuntimeError(error_msg)
    
    def _run_task(self, file_path: str) -> str:
        """
        Internal method to run the task in a separate thread
        
        Args:
            file_path: Path to the uploaded file
            
        Returns:
            Path to the configuration file
        """
        try:
            # Use the default yourbench organization
            org_name = os.getenv("HF_ORGANIZATION")
            
            # Check if HF token is available
            hf_token = os.getenv("HF_TOKEN")
            if not hf_token:
                raise RuntimeError("HF_TOKEN environment variable is not defined")
            
            self._add_log(f"[INFO] Organization: {org_name}")
            
            time.sleep(0.5)  # Simulate delay
            
            # Save the uploaded file
            saved_file_path = self.save_uploaded_file(file_path)
            
            time.sleep(1)  # Simulate delay
            
            # Path for the config file
            config_dir = pathlib.Path(f"uploaded_files/{self.session_uid}")
            config_path = config_dir / "config.yml"
            
            # Generate dataset name based on session ID
            dataset_name = f"yourbench_{self.session_uid}"
            self._add_log(f"[INFO] Dataset name: {dataset_name}")
            
            time.sleep(0.8)  # Simulate delay
            
            # Generate and save the configuration
            config = self.generate_base_config(org_name, dataset_name)
            
            time.sleep(1.2)  # Simulate delay
            
            config_file_path = self.save_yaml_file(config, str(config_path))
            
            self._add_log(f"[INFO] Configuration generated successfully: {config_file_path}")
            
            # Simulate additional processing
            time.sleep(1.5)  # Simulate delay
            self._add_log("[INFO] Starting ingestion")
            
            time.sleep(2)  # Simulate delay
            self._add_log(f"[INFO] Processing file: {dataset_name}")
            
            time.sleep(2)  # Simulate delay
            self._add_log("[SUCCESS] Stage completed: config_generation")
            
            # Tâche terminée
            self.mark_task_completed()
            
            return str(config_path)
        except Exception as e:
            error_msg = f"Error generating configuration: {str(e)}"
            self._add_log(f"[ERROR] {error_msg}")
            self.mark_task_completed()
            raise RuntimeError(error_msg)
    
    def run(self, file_path: str, token: Optional[str] = None) -> str:
        """
        Run the task to create and save the configuration file asynchronously
        
        Args:
            file_path: Path to the uploaded file
            token: Hugging Face token (not used, using HF_TOKEN from environment)
            
        Returns:
            Path to the configuration file
        """
        # Mark the task as running
        self.is_running_flag.set()
        
        # Start the task in a separate thread
        self.thread = threading.Thread(target=self._run_task, args=(file_path,))
        self.thread.daemon = True
        self.thread.start()
        
        # Return the expected config path
        return f"uploaded_files/{self.session_uid}/config.yml"
    
    def is_running(self) -> bool:
        """
        Check if the task is running
        
        Returns:
            True if running, False otherwise
        """
        return self.is_running_flag.is_set() and not self.is_completed
    
    def is_task_completed(self) -> bool:
        """
        Check if the task is completed
        
        Returns:
            True if completed, False otherwise
        """
        return self.is_completed
    
    def mark_task_completed(self) -> None:
        """
        Mark the task as completed
        """
        self.is_completed = True
        self.is_running_flag.clear()
        self._add_log("[INFO] Configuration generation task completed")