backend/tasks/createBench.py

"""
Task to ingest and transform documents to markdown using yourbench
"""
import os
import time
import pathlib
import subprocess
import threading
from typing import Optional, List, Tuple, Dict, Any
import yaml

from loguru import logger


class CreateBenchTask:
    """
    Task to ingest and transform documents to markdown using yourbench
    """

    def __init__(self, session_uid: str, config_path: Optional[str] = None):
        """
        Initialize the ingestion task
        
        Args:
            session_uid: Session ID for this task
            config_path: Path to the configuration file, will be generated if None
        """
        self.session_uid = session_uid
        self.logs: List[str] = []
        self.is_completed = False
        self.process = None
        self.is_running_flag = threading.Event()
        
        # Default config path if not provided
        if config_path is None:
            config_path = f"uploaded_files/{session_uid}/config.yml"
        self.config_path = config_path
        
        # Command to run yourbench - modified to avoid error with uv run
        self.command = ["yourbench", "run", "--config", str(self.config_path)]
        
        self._add_log("[INFO] Initializing ingestion task")
        self._add_log(f"[INFO] Using configuration file: {self.config_path}")
    
    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 copy of the list to avoid reference problems
            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()  # Return a copy to avoid reference problems
    
    def is_task_completed(self) -> bool:
        """
        Check if the task is completed
        
        Returns:
            True if completed, False otherwise
        """
        return self.is_completed
    
    def is_running(self) -> bool:
        """
        Check if the process is running
        
        Returns:
            True if running, False otherwise
        """
        return self.is_running_flag.is_set()
    
    def stop(self) -> None:
        """
        Stop the process if it's running
        """
        if self.process and self.is_running():
            self._add_log("[INFO] Stopping ingestion process")
            try:
                self.process.terminate()
                # Wait 5 seconds for termination
                self.process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self._add_log("[WARN] Process not responding, forcing termination")
                self.process.kill()
            finally:
                self.is_running_flag.clear()
                self.is_completed = True
                self._add_log("[INFO] Ingestion process stopped")
    
    def _capture_output(self) -> None:
        """
        Capture and process the output from the yourbench process
        """
        self._add_log("[INFO] Starting output capture")
        
        try:
            while self.is_running() and self.process:
                line = self.process.stdout.readline()
                if not line:
                    # If no line is read and the process is no longer running
                    if self.process.poll() is not None:
                        self.is_running_flag.clear()
                        break
                    # Otherwise, wait a bit and continue
                    time.sleep(0.1)
                    continue
                
                # Process the output line
                line = line.strip()
                if line:
                    # Log raw output for debugging
                    self._add_log(f"[DEBUG] Raw output: {line}")
                    # Filter and format the line as needed
                    if "ERROR" in line:
                        self._add_log(f"[ERROR] {line}")
                    elif "WARNING" in line:
                        self._add_log(f"[WARN] {line}")
                    else:
                        # Detect completed stages
                        if "Completed stage:" in line:
                            stage = line.split("'")[1] if "'" in line else line
                            self._add_log(f"[SUCCESS] Stage completed: {stage}")
                        else:
                            self._add_log(f"[INFO] {line}")
            
            # Check exit code once the process is finished
            if self.process:
                exit_code = self.process.poll()
                if exit_code == 0:
                    self._add_log("[SUCCESS] Ingestion process completed successfully")
                else:
                    self._add_log(f"[ERROR] Ingestion process terminated with error code: {exit_code}")
        except Exception as e:
            self._add_log(f"[ERROR] Error during output capture: {str(e)}")
        finally:
            self.is_completed = True
            self.is_running_flag.clear()
            self._add_log("[INFO] Output capture completed")
    
    def run(self, token: Optional[str] = None) -> None:
        """
        Run the ingestion task
        
        Args:
            token: Hugging Face token
        """
        try:
            self._add_log("[INFO] Starting ingestion process")
            
            # Check if the configuration file exists
            if not os.path.exists(self.config_path):
                raise FileNotFoundError(f"Configuration file does not exist: {self.config_path}")
            
            # Examine the configuration to get information
            try:
                with open(self.config_path, 'r') as f:
                    config_yaml = yaml.safe_load(f)
                    
                # Get source and destination paths
                source_dir = config_yaml.get("pipeline", {}).get("ingestion", {}).get("source_documents_dir", "")
                output_dir = config_yaml.get("pipeline", {}).get("ingestion", {}).get("output_dir", "")
                
                if source_dir:
                    self._add_log(f"[INFO] Source directory: {source_dir}")
                if output_dir:
                    self._add_log(f"[INFO] Output directory: {output_dir}")
                    
                # List files to process if the directory exists
                if source_dir and os.path.exists(source_dir):
                    files = os.listdir(source_dir)
                    if files:
                        self._add_log(f"[INFO] Files to process: {', '.join(files)}")
                    else:
                        self._add_log("[WARN] No files found in source directory")
                
            except Exception as e:
                self._add_log(f"[WARN] Unable to read configuration: {str(e)}")
            
            # Environment preparation
            env = os.environ.copy()
            
            # Explicitly define environment variables for authentication
            hf_token = os.getenv("HF_TOKEN")
            if hf_token:
                # Explicitly export these variables for yourbench
                env["HF_TOKEN"] = hf_token
                env["HUGGING_FACE_HUB_TOKEN"] = hf_token
                env["HF_ORGANIZATION"] = os.getenv("HF_ORGANIZATION", "yourbench")
                self._add_log("[INFO] Environment variables HF_TOKEN, HUGGING_FACE_HUB_TOKEN and HF_ORGANIZATION exported")
            
            # In development mode, only simulate ingestion
            if os.environ.get("DEVELOPMENT_MODE", "").lower() == "true":
                self._add_log("[INFO] Development mode enabled, simulating ingestion")
                self._simulate_ingestion_process()
                return
            
            # Start the process
            self._add_log(f"[INFO] Executing command: {' '.join(self.command)}")
            
            self.process = subprocess.Popen(
                self.command,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
                universal_newlines=True,
                env=env
            )
            
            # Mark the process as running
            self.is_running_flag.set()
            
            # Start a thread to capture output
            output_thread = threading.Thread(target=self._capture_output)
            output_thread.daemon = True
            output_thread.start()
            
            self._add_log(f"[INFO] Process started with PID: {self.process.pid}")
            
        except Exception as e:
            self._add_log(f"[ERROR] Error starting ingestion process: {str(e)}")
            self.is_completed = True
    
    def _simulate_ingestion_process(self) -> None:
        """
        Simulate the ingestion process for testing/development
        This will be removed in production
        """
        # This method is just to simulate logs during development
        # It will be removed in production
        
        threading.Thread(target=self._simulate_logs).start()
    
    def _simulate_logs(self) -> None:
        """
        Simulate logs for testing/development
        This will be used when yourbench isn't installed or in development mode
        """
        # Log simulation (used when yourbench is not available)
        self._add_log("[INFO] Simulation mode enabled (yourbench is not actually running)")
        
        # Get filenames from source directory
        source_files = []
        try:
            with open(self.config_path, 'r') as f:
                config_yaml = yaml.safe_load(f)
                
            source_dir = config_yaml.get("pipeline", {}).get("ingestion", {}).get("source_documents_dir", "")
            if source_dir and os.path.exists(source_dir):
                source_files = [f for f in os.listdir(source_dir) 
                               if os.path.isfile(os.path.join(source_dir, f))]
        except Exception:
            source_files = ["document.pdf", "document.txt"]  # Fallback
        
        # Create output directory if it doesn't exist
        output_dir = ""
        try:
            output_dir = config_yaml.get("pipeline", {}).get("ingestion", {}).get("output_dir", "")
            if output_dir:
                os.makedirs(output_dir, exist_ok=True)
        except Exception:
            pass
            
        # Simulate file processing
        time.sleep(1)
        self._add_log("[INFO] Initializing document ingestion")
        time.sleep(1.5)
        self._add_log("[INFO] Loading configuration parameters")
        time.sleep(1)
        self._add_log("[INFO] Verifying source files")
        
        # Process each file
        for file in source_files:
            time.sleep(1.5)
            self._add_log(f"[INFO] Processing file: {file}")
            time.sleep(2)
            self._add_log(f"[INFO] Extracting content from {file}")
            time.sleep(1.5)
            self._add_log(f"[INFO] Converting to markdown: {file}")
            
            # Create a simulated markdown file if an output directory is defined
            if output_dir:
                base_name = os.path.splitext(file)[0]
                output_file = os.path.join(output_dir, f"{base_name}.md")
                try:
                    with open(output_file, 'w') as f:
                        f.write(f"# {base_name}\n\n")
                        f.write("This is a markdown document automatically generated by the simulation.\n\n")
                        f.write("## Section 1\n\n")
                        f.write("Content of section 1...\n\n")
                        f.write("## Section 2\n\n")
                        f.write("Content of section 2...\n\n")
                    self._add_log(f"[INFO] Markdown file created: {output_file}")
                except Exception as e:
                    self._add_log(f"[ERROR] Error creating markdown file: {str(e)}")
        
        time.sleep(2)
        self._add_log("[INFO] Finalizing processing")
        time.sleep(1)
        self._add_log("[SUCCESS] Stage completed: ingestion")
        time.sleep(0.5)
        self._add_log("[SUCCESS] Ingestion completed successfully")
        
        # Mark task as completed
        self.is_completed = True